{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[],"gpuType":"L4","collapsed_sections":["yXwXG1r0yP__","v8Fimlw_yZj_","-Vj68rbBEKES","VV4alT_R22se","TA4DWJ3Qn60x","UQxAwmUYnWkp","SjjX6nDgnKwI","Ktwbrdh7m_gv","Ce0Ef-gOobJn"],"machine_shape":"hm","mount_file_id":"1QYHy-W18FL_SXX_AIKOorvypjQoeP8XT","authorship_tag":"ABX9TyO56fNtmTGlHAqbAIMZom4m"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"},"accelerator":"GPU"},"cells":[{"cell_type":"code","execution_count":1,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"djpBLPTRl2fW","executionInfo":{"status":"ok","timestamp":1721025358263,"user_tz":-480,"elapsed":1273,"user":{"displayName":"知否","userId":"12758421836332615961"}},"outputId":"a9e17d46-8547-43ee-ad63-375a1f04a651"},"outputs":[{"output_type":"stream","name":"stdout","text":["Python 3.10.12\n"]}],"source":["!python -V\n"]},{"cell_type":"markdown","source":["# 导入所需要环境"],"metadata":{"id":"yXwXG1r0yP__"}},{"cell_type":"code","source":["!pip3 install einops matplotlib matplotlib scikit_learn scipy scipy open3d meshio"],"metadata":{"id":"76-vyRjomPMh","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1721025384886,"user_tz":-480,"elapsed":26012,"user":{"displayName":"知否","userId":"12758421836332615961"}},"outputId":"2407269f-5324-40f0-b1ec-2e431df10e9d"},"execution_count":2,"outputs":[{"output_type":"stream","name":"stdout","text":["Collecting einops\n","  Downloading einops-0.8.0-py3-none-any.whl (43 kB)\n","\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/43.2 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m43.2/43.2 kB\u001b[0m \u001b[31m2.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[?25hRequirement already satisfied: matplotlib in /usr/local/lib/python3.10/dist-packages (3.7.1)\n","Requirement already satisfied: scikit_learn in /usr/local/lib/python3.10/dist-packages (1.2.2)\n","Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (1.11.4)\n","Collecting open3d\n","  Downloading open3d-0.18.0-cp310-cp310-manylinux_2_27_x86_64.whl (399.7 MB)\n","\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m399.7/399.7 MB\u001b[0m \u001b[31m1.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[?25hCollecting meshio\n","  Downloading meshio-5.3.5-py3-none-any.whl (166 kB)\n","\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m166.2/166.2 kB\u001b[0m \u001b[31m25.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[?25hRequirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib) (1.2.1)\n","Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib) (0.12.1)\n","Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib) (4.53.1)\n","Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib) (1.4.5)\n","Requirement already satisfied: numpy>=1.20 in /usr/local/lib/python3.10/dist-packages (from matplotlib) (1.25.2)\n","Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib) (24.1)\n","Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib) (9.4.0)\n","Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib) (3.1.2)\n","Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib) (2.8.2)\n","Requirement already satisfied: joblib>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from scikit_learn) (1.4.2)\n","Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from scikit_learn) (3.5.0)\n","Collecting dash>=2.6.0 (from open3d)\n","  Downloading dash-2.17.1-py3-none-any.whl (7.5 MB)\n","\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.5/7.5 MB\u001b[0m \u001b[31m103.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[?25hRequirement already satisfied: werkzeug>=2.2.3 in /usr/local/lib/python3.10/dist-packages (from open3d) (3.0.3)\n","Requirement already satisfied: nbformat>=5.7.0 in /usr/local/lib/python3.10/dist-packages (from open3d) (5.10.4)\n","Collecting configargparse (from open3d)\n","  Downloading ConfigArgParse-1.7-py3-none-any.whl (25 kB)\n","Collecting ipywidgets>=8.0.4 (from open3d)\n","  Downloading ipywidgets-8.1.3-py3-none-any.whl (139 kB)\n","\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m139.4/139.4 kB\u001b[0m \u001b[31m23.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[?25hCollecting addict (from open3d)\n","  Downloading addict-2.4.0-py3-none-any.whl (3.8 kB)\n","Requirement already satisfied: pandas>=1.0 in /usr/local/lib/python3.10/dist-packages (from open3d) (2.0.3)\n","Requirement already satisfied: pyyaml>=5.4.1 in /usr/local/lib/python3.10/dist-packages (from open3d) (6.0.1)\n","Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from open3d) (4.66.4)\n","Collecting pyquaternion (from open3d)\n","  Downloading pyquaternion-0.9.9-py3-none-any.whl (14 kB)\n","Requirement already satisfied: rich in /usr/local/lib/python3.10/dist-packages (from meshio) (13.7.1)\n","Requirement already satisfied: Flask<3.1,>=1.0.4 in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (2.2.5)\n","Requirement already satisfied: plotly>=5.0.0 in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (5.15.0)\n","Collecting dash-html-components==2.0.0 (from dash>=2.6.0->open3d)\n","  Downloading dash_html_components-2.0.0-py3-none-any.whl (4.1 kB)\n","Collecting dash-core-components==2.0.0 (from dash>=2.6.0->open3d)\n","  Downloading dash_core_components-2.0.0-py3-none-any.whl (3.8 kB)\n","Collecting dash-table==5.0.0 (from dash>=2.6.0->open3d)\n","  Downloading dash_table-5.0.0-py3-none-any.whl (3.9 kB)\n","Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (8.0.0)\n","Requirement already satisfied: typing-extensions>=4.1.1 in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (4.12.2)\n","Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (2.31.0)\n","Collecting retrying (from dash>=2.6.0->open3d)\n","  Downloading retrying-1.3.4-py3-none-any.whl (11 kB)\n","Requirement already satisfied: nest-asyncio in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (1.6.0)\n","Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (67.7.2)\n","Collecting comm>=0.1.3 (from ipywidgets>=8.0.4->open3d)\n","  Downloading comm-0.2.2-py3-none-any.whl (7.2 kB)\n","Requirement already satisfied: ipython>=6.1.0 in /usr/local/lib/python3.10/dist-packages (from ipywidgets>=8.0.4->open3d) (7.34.0)\n","Requirement already satisfied: traitlets>=4.3.1 in /usr/local/lib/python3.10/dist-packages (from ipywidgets>=8.0.4->open3d) (5.7.1)\n","Collecting widgetsnbextension~=4.0.11 (from ipywidgets>=8.0.4->open3d)\n","  Downloading widgetsnbextension-4.0.11-py3-none-any.whl (2.3 MB)\n","\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.3/2.3 MB\u001b[0m \u001b[31m94.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[?25hRequirement already satisfied: jupyterlab-widgets~=3.0.11 in /usr/local/lib/python3.10/dist-packages (from ipywidgets>=8.0.4->open3d) (3.0.11)\n","Requirement already satisfied: fastjsonschema>=2.15 in /usr/local/lib/python3.10/dist-packages (from nbformat>=5.7.0->open3d) (2.20.0)\n","Requirement already satisfied: jsonschema>=2.6 in /usr/local/lib/python3.10/dist-packages (from nbformat>=5.7.0->open3d) (4.19.2)\n","Requirement already satisfied: jupyter-core!=5.0.*,>=4.12 in /usr/local/lib/python3.10/dist-packages (from nbformat>=5.7.0->open3d) (5.7.2)\n","Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.0->open3d) (2023.4)\n","Requirement already satisfied: tzdata>=2022.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.0->open3d) (2024.1)\n","Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.7->matplotlib) (1.16.0)\n","Requirement already satisfied: MarkupSafe>=2.1.1 in /usr/local/lib/python3.10/dist-packages (from werkzeug>=2.2.3->open3d) (2.1.5)\n","Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.10/dist-packages (from rich->meshio) (3.0.0)\n","Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.10/dist-packages (from rich->meshio) (2.16.1)\n","Requirement already satisfied: Jinja2>=3.0 in /usr/local/lib/python3.10/dist-packages (from Flask<3.1,>=1.0.4->dash>=2.6.0->open3d) (3.1.4)\n","Requirement already satisfied: itsdangerous>=2.0 in /usr/local/lib/python3.10/dist-packages (from Flask<3.1,>=1.0.4->dash>=2.6.0->open3d) (2.2.0)\n","Requirement already satisfied: click>=8.0 in /usr/local/lib/python3.10/dist-packages (from Flask<3.1,>=1.0.4->dash>=2.6.0->open3d) (8.1.7)\n","Collecting jedi>=0.16 (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d)\n","  Downloading jedi-0.19.1-py2.py3-none-any.whl (1.6 MB)\n","\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m87.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n","\u001b[?25hRequirement already satisfied: decorator in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (4.4.2)\n","Requirement already satisfied: pickleshare in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (0.7.5)\n","Requirement already satisfied: prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (3.0.47)\n","Requirement already satisfied: backcall in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (0.2.0)\n","Requirement already satisfied: matplotlib-inline in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (0.1.7)\n","Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (4.9.0)\n","Requirement already satisfied: attrs>=22.2.0 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=2.6->nbformat>=5.7.0->open3d) (23.2.0)\n","Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=2.6->nbformat>=5.7.0->open3d) (2023.12.1)\n","Requirement already satisfied: referencing>=0.28.4 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=2.6->nbformat>=5.7.0->open3d) (0.35.1)\n","Requirement already satisfied: rpds-py>=0.7.1 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=2.6->nbformat>=5.7.0->open3d) (0.19.0)\n","Requirement already satisfied: platformdirs>=2.5 in /usr/local/lib/python3.10/dist-packages (from jupyter-core!=5.0.*,>=4.12->nbformat>=5.7.0->open3d) (4.2.2)\n","Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.10/dist-packages (from markdown-it-py>=2.2.0->rich->meshio) (0.1.2)\n","Requirement already satisfied: tenacity>=6.2.0 in /usr/local/lib/python3.10/dist-packages (from plotly>=5.0.0->dash>=2.6.0->open3d) (8.5.0)\n","Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.10/dist-packages (from importlib-metadata->dash>=2.6.0->open3d) (3.19.2)\n","Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->dash>=2.6.0->open3d) (3.3.2)\n","Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->dash>=2.6.0->open3d) (3.7)\n","Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->dash>=2.6.0->open3d) (2.0.7)\n","Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->dash>=2.6.0->open3d) (2024.7.4)\n","Requirement already satisfied: parso<0.9.0,>=0.8.3 in /usr/local/lib/python3.10/dist-packages (from jedi>=0.16->ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (0.8.4)\n","Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.10/dist-packages (from pexpect>4.3->ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (0.7.0)\n","Requirement already satisfied: wcwidth in /usr/local/lib/python3.10/dist-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (0.2.13)\n","Installing collected packages: dash-table, dash-html-components, dash-core-components, addict, widgetsnbextension, retrying, pyquaternion, jedi, einops, configargparse, comm, meshio, ipywidgets, dash, open3d\n","  Attempting uninstall: widgetsnbextension\n","    Found existing installation: widgetsnbextension 3.6.7\n","    Uninstalling widgetsnbextension-3.6.7:\n","      Successfully uninstalled widgetsnbextension-3.6.7\n","  Attempting uninstall: ipywidgets\n","    Found existing installation: ipywidgets 7.7.1\n","    Uninstalling ipywidgets-7.7.1:\n","      Successfully uninstalled ipywidgets-7.7.1\n","Successfully installed addict-2.4.0 comm-0.2.2 configargparse-1.7 dash-2.17.1 dash-core-components-2.0.0 dash-html-components-2.0.0 dash-table-5.0.0 einops-0.8.0 ipywidgets-8.1.3 jedi-0.19.1 meshio-5.3.5 open3d-0.18.0 pyquaternion-0.9.9 retrying-1.3.4 widgetsnbextension-4.0.11\n"]}]},{"cell_type":"markdown","source":["# 下载官方数据"],"metadata":{"id":"v8Fimlw_yZj_"}},{"cell_type":"code","source":["# !wget --header=\"Host: drive.usercontent.google.com\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\" --header=\"Accept-Language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6\" --header=\"Cookie: __Secure-ENID=12.SE=Yd0Bj-CLJ14fnd4qzdJHmwUs4B5zz46UaPC1cPJigNqqFV9PtM2CYyBpSbCkOyzUwzlEdZ1nZFf-igtGi7wSdJ_gqQSfQfh84r9egqFQAy9-GKayCRbdQKdera-2mkpuIT-c64CyR9vfNojM3hxZ9Dej-dGvtxlGjal9ttEHybw; __gsas=ID=ae0421b9a34b478c:T=1710758437:RT=1710758437:S=ALNI_MZP13R9ZOHbCzC0rgHSMrGXj6GCsg; HSID=A-4I-ZudDNUIB6EKH; SSID=A7v_1v9un6xAwVNku; APISID=ctK8IbLjeuDUmgys/AFnMSLWt9KddceDI6; SAPISID=J7GhTwED67EBqJJT/A9nwK7mr0ijGPw08r; __Secure-1PAPISID=J7GhTwED67EBqJJT/A9nwK7mr0ijGPw08r; __Secure-3PAPISID=J7GhTwED67EBqJJT/A9nwK7mr0ijGPw08r; SID=g.a000kgiBabgKCiCYKve9zfoWVgz9eu8sBA6N4XDPPpP5pcW16_C_kzuBV1TvOhAIC8VF1e9fpgACgYKATQSARQSFQHGX2Mi8LXUwWoIwNCEPU8Sy3mXUxoVAUF8yKqGXVfjTGz9gQal7nwGr4Pl0076; __Secure-1PSID=g.a000kgiBabgKCiCYKve9zfoWVgz9eu8sBA6N4XDPPpP5pcW16_C_PDa-DzVmbdGFPyxMQpk9_QACgYKAewSARQSFQHGX2MiAeee4fn0OWglWZfAygqkyBoVAUF8yKp-Sfmtnueimxc-0QbJRF9I0076; __Secure-3PSID=g.a000kgiBabgKCiCYKve9zfoWVgz9eu8sBA6N4XDPPpP5pcW16_C_g9IrMeU98APBo9Stp6wEnAACgYKAQASARQSFQHGX2MiFWtc9ucONXnpxBzlRdudEhoVAUF8yKoeZwCpJDnjfAFjGssHSUGm0076; NID=515=GQhY9nKKFCx3qFDjE0MA4ubjWNdef6xCIY_RfWOPWKEtyfBN3nAUl8WHI2VczjNQ4rVkj1XBAY8WNWHXyqSK10CfT4FxsFlPzrHIJpeTtm1nWRNBd9AAfBKJHz4XpESszntVUTE_59RklZuKo0vk1poReVi2da1PZKC3CTKH2Ll3gB5xuB9wf4bmq8ylVUuIROPJczr0XnCuUHV3qLdBvgy9_870b6UwOq1iOlIxFQFm01EZ4pqF4q1Ub3QRSWpEMLh4LSZFpJ5O255R5OV7krmEdDvH_sHoTEPZAg2PoEpwAyGK6Xp9qcLIlldgx5-5V86N8Wtb93uTlQuA_CFXb5_2eP3bgeX8txwlJ5SrldVjg9ctzYtBU2RwJKTSvdHfIG7lpOkg6XlkvDOcJpR3DihT_OlqnPn7drCAJpvVDv29hZn5XPMXaSrNdbG64OJ9urJEw5odEwsLYkkpC1vmlUcuoo52S5f6RQu0Z8kZiV8iRW6XIqHsSmQHunVaxk6xWCStUg; __Secure-1PSIDTS=sidts-CjEB3EgAEtTS0OazynCofIH4RCBstiRP5flEcvYW3z4Fg9oGd5QOESDOZt1wO2iqUYHjEAA; __Secure-3PSIDTS=sidts-CjEB3EgAEtTS0OazynCofIH4RCBstiRP5flEcvYW3z4Fg9oGd5QOESDOZt1wO2iqUYHjEAA; SIDCC=AKEyXzVI6aMX8lSDja86Yts3FBAtBzPCzVNgaX5BCz78NWsWzlT3yFWKUV7ZE46SFzE1GiBI-cHdTw; __Secure-1PSIDCC=AKEyXzUo4NQAwqqPMxP2eye-MFEbZmBIm_sZqRU1amttg0YoQkc8ZKSNXdHl5jNCMEbhrUHhS9-K; __Secure-3PSIDCC=AKEyXzWf2lIdmDLeZKpXSi9GytVQb6XudrYiNUBA5gW952YuLh8kL6T3IbBlu8zOTfGEcdUp5O1R\" --header=\"Connection: keep-alive\" \"https://drive.usercontent.google.com/download?id=1JwR0Q1ArTg6c47EF2ZuIBpQwCPgXKrO2&export=download&authuser=0&confirm=t&uuid=dc3aa13c-c3a9-458f-983a-8586798cb635&at=APZUnTX25XMxi-z-3wBcgR93IGsL%3A1719235792953\" -c -O 'Dataset.zip'\n","# !unzip Dataset.zip"],"metadata":{"id":"-iCVuuqnQACM","executionInfo":{"status":"ok","timestamp":1721025384887,"user_tz":-480,"elapsed":7,"user":{"displayName":"知否","userId":"12758421836332615961"}}},"execution_count":3,"outputs":[]},{"cell_type":"code","source":["\n","!wget --header=\"Host: ai-studio-online.bj.bcebos.com\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36 Edg/126.0.0.0\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\" --header=\"Accept-Language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6\" --header=\"Referer: https://aistudio.baidu.com/\" \"https://ai-studio-online.bj.bcebos.com/v1/1638f9c292b9437bb46885186407a63e584856c91f9f4c18908b87abd46471e0?responseContentDisposition=attachment%3B%20filename%3Dtrack_B.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-05-05T03%3A02%3A25Z%2F-1%2F%2Fcfdfd6b6a9e096c761ee8e7d863d586741c69a9e6de89f9c3696706d35f8b265\" -c -O 'track_B.zip'\n","!wget --header=\"Host: ai-studio-online.bj.bcebos.com\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36 Edg/126.0.0.0\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\" --header=\"Accept-Language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6\" --header=\"Referer: https://aistudio.baidu.com/\" \"https://ai-studio-online.bj.bcebos.com/v1/2dddd05e577849ad95e1fe1133d3af29d13085ac0cfd499c853ff5d9df2ac07f?responseContentDisposition=attachment%3B%20filename%3Dtrain_data.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-05-05T03%3A23%3A33Z%2F-1%2F%2F8540633c7e39fddf8471d6d8206c3b761748c58c06005acb218593a8df19d7f1\" -c -O 'train_data.zip'\n","!wget --header=\"Host: ai-studio-online.bj.bcebos.com\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36 Edg/126.0.0.0\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\" --header=\"Accept-Language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6\" --header=\"Referer: https://aistudio.baidu.com/\" \"https://ai-studio-online.bj.bcebos.com/v1/a96dc8ba8201445b966980a0a48f52705338a48e29e64c53bddb7ef8861c5123?responseContentDisposition=attachment%3B%20filename%3Dtrack_A.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-05-06T07%3A54%3A51Z%2F-1%2F%2F17b5155bd16a8af1e4af971498082687656af7fcecfc5a8e57591b85053210ec\" -c -O 'track_A.zip'\n","# !wget --header=\"Host: ai-studio-online.bj.bcebos.com\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36 Edg/126.0.0.0\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\" --header=\"Accept-Language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6\" --header=\"Referer: https://aistudio.baidu.com/\" \"https://ai-studio-online.bj.bcebos.com/v1/dcaba9f0e87549e395e1682a4a0a43c547a896034e3f4417a4d10ba85a949944?responseContentDisposition=attachment%3B%20filename%3DPaddleScience.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-05-11T16%3A01%3A57Z%2F-1%2F%2Fb2149791689d3b19c02a86782229901cb75743d07c7163dfa185fea286cd1f02\" -c -O 'PaddleScience.zip'\n","!wget --header=\"Host: ai-studio-online.bj.bcebos.com\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36 Edg/126.0.0.0\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\" --header=\"Accept-Language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6\" --header=\"Referer: https://aistudio.baidu.com/\" \"https://ai-studio-online.bj.bcebos.com/v1/38e9adf0fce84527aad3558cc3e82d0e9a251aac4c934297afae9b74d9b3d1e9?responseContentDisposition=attachment%3B%20filename%3Dtrain_track_B.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-06-04T03%3A21%3A02Z%2F-1%2F%2Facd359add161bace603a52c7a268467406cb3c1889a7114bbb687de8002b55f6\" -c -O 'train_track_B.zip'\n","!wget --header=\"Host: ai-studio-online.bj.bcebos.com\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36 Edg/126.0.0.0\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\" --header=\"Accept-Language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6\" --header=\"Referer: https://aistudio.baidu.com/\" \"https://ai-studio-online.bj.bcebos.com/v1/7877c2fd300345599ed3365feda50425f4857caa71bf4af9bf047fb08e35aa97?responseContentDisposition=attachment%3B%20filename%3Dmesh_B_0603.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-06-04T11%3A48%3A26Z%2F-1%2F%2Fec6cce492ba6afff841ef197860065742a7ba8220def02c1cbc311333aa993b2\" -c -O 'mesh_B_0603.zip'\n","# !wget --header=\"Host: ai-studio-online.bj.bcebos.com\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36 Edg/126.0.0.0\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\" --header=\"Accept-Language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6\" --header=\"Referer: https://aistudio.baidu.com/\" \"https://ai-studio-online.bj.bcebos.com/v1/a02dba5700974c6a811f579fff216ccf9a4129b849994dcca99390b222c28572?responseContentDisposition=attachment%3B%20filename%3D3rd_lib.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-06-08T05%3A29%3A20Z%2F-1%2F%2Fef71819149664a3c8438d3dfd02544e77cdb4fda679efbdd402b4e2c77a2c06d\" -c -O '3rd_lib.zip'"],"metadata":{"id":"5mbGmjSnQqYJ","colab":{"base_uri":"https://localhost:8080/"},"outputId":"cb56e166-66d6-4e30-f2fd-152d4302e3cc"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["--2024-07-15 06:36:24--  https://ai-studio-online.bj.bcebos.com/v1/1638f9c292b9437bb46885186407a63e584856c91f9f4c18908b87abd46471e0?responseContentDisposition=attachment%3B%20filename%3Dtrack_B.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-05-05T03%3A02%3A25Z%2F-1%2F%2Fcfdfd6b6a9e096c761ee8e7d863d586741c69a9e6de89f9c3696706d35f8b265\n","Resolving ai-studio-online.bj.bcebos.com (ai-studio-online.bj.bcebos.com)... 103.235.47.176, 2409:8c04:1001:1203:0:ff:b0bb:4f27\n","Connecting to ai-studio-online.bj.bcebos.com (ai-studio-online.bj.bcebos.com)|103.235.47.176|:443... connected.\n","HTTP request sent, awaiting response... 200 OK\n","Length: 1012191818 (965M) [application/octet-stream]\n","Saving to: ‘track_B.zip’\n","\n","track_B.zip         100%[===================>] 965.30M  15.8MB/s    in 79s     \n","\n","2024-07-15 06:37:44 (12.3 MB/s) - ‘track_B.zip’ saved [1012191818/1012191818]\n","\n","--2024-07-15 06:37:44--  https://ai-studio-online.bj.bcebos.com/v1/2dddd05e577849ad95e1fe1133d3af29d13085ac0cfd499c853ff5d9df2ac07f?responseContentDisposition=attachment%3B%20filename%3Dtrain_data.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-05-05T03%3A23%3A33Z%2F-1%2F%2F8540633c7e39fddf8471d6d8206c3b761748c58c06005acb218593a8df19d7f1\n","Resolving ai-studio-online.bj.bcebos.com (ai-studio-online.bj.bcebos.com)... 103.235.47.176, 2409:8c04:1001:1203:0:ff:b0bb:4f27\n","Connecting to ai-studio-online.bj.bcebos.com (ai-studio-online.bj.bcebos.com)|103.235.47.176|:443... connected.\n","HTTP request sent, awaiting response... 200 OK\n","Length: 46031310 (44M) [application/octet-stream]\n","Saving to: ‘train_data.zip’\n","\n","train_data.zip      100%[===================>]  43.90M  12.1MB/s    in 18s     \n","\n","2024-07-15 06:38:04 (2.39 MB/s) - ‘train_data.zip’ saved [46031310/46031310]\n","\n","--2024-07-15 06:38:04--  https://ai-studio-online.bj.bcebos.com/v1/a96dc8ba8201445b966980a0a48f52705338a48e29e64c53bddb7ef8861c5123?responseContentDisposition=attachment%3B%20filename%3Dtrack_A.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-05-06T07%3A54%3A51Z%2F-1%2F%2F17b5155bd16a8af1e4af971498082687656af7fcecfc5a8e57591b85053210ec\n","Resolving ai-studio-online.bj.bcebos.com (ai-studio-online.bj.bcebos.com)... 103.235.47.176, 2409:8c04:1001:1203:0:ff:b0bb:4f27\n","Connecting to ai-studio-online.bj.bcebos.com (ai-studio-online.bj.bcebos.com)|103.235.47.176|:443... connected.\n","HTTP request sent, awaiting response... 200 OK\n","Length: 4688102 (4.5M) [application/octet-stream]\n","Saving to: ‘track_A.zip’\n","\n","track_A.zip         100%[===================>]   4.47M   940KB/s    in 14s     \n","\n","2024-07-15 06:38:20 (316 KB/s) - ‘track_A.zip’ saved [4688102/4688102]\n","\n","--2024-07-15 06:38:20--  https://ai-studio-online.bj.bcebos.com/v1/38e9adf0fce84527aad3558cc3e82d0e9a251aac4c934297afae9b74d9b3d1e9?responseContentDisposition=attachment%3B%20filename%3Dtrain_track_B.zip&authorization=bce-auth-v1%2F5cfe9a5e1454405eb2a975c43eace6ec%2F2024-06-04T03%3A21%3A02Z%2F-1%2F%2Facd359add161bace603a52c7a268467406cb3c1889a7114bbb687de8002b55f6\n","Resolving ai-studio-online.bj.bcebos.com (ai-studio-online.bj.bcebos.com)... 103.235.47.176, 2409:8c04:1001:1203:0:ff:b0bb:4f27\n","Connecting to ai-studio-online.bj.bcebos.com (ai-studio-online.bj.bcebos.com)|103.235.47.176|:443... connected.\n","HTTP request sent, awaiting response... 200 OK\n","Length: 4740031429 (4.4G) [application/octet-stream]\n","Saving to: ‘train_track_B.zip’\n","\n","train_track_B.zip    22%[===>                ]   1004M  20.2MB/s    eta 3m 52s "]}]},{"cell_type":"code","source":["!mkdir Dataset\n","!unzip -o track_B.zip -d Dataset/\n","!unzip -o train_data.zip -d Dataset/\n","!unzip -o track_A.zip -d Dataset/\n","!mkdir Dataset/train_track_B && unzip -o train_track_B.zip -d Dataset/train_track_B/\n","!unzip -o mesh_B_0603.zip -d Dataset/\n"],"metadata":{"id":"d0-rRJVfQ5Ge"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# 数据处理"],"metadata":{"id":"-Vj68rbBEKES"}},{"cell_type":"markdown","source":["以下代码是将数据处理成官方文件夹的形式，单纯文件移动操作，不引入新的数据"],"metadata":{"id":"Z9z9-NNU3F46"}},{"cell_type":"code","source":["import os\n","import shutil\n","\n","# 指定 train_data 目录路径\n","path = 'Dataset/data'\n","\n","# 新建 Feature 和 Label 文件夹\n","Feature_path = os.path.join(path, 'Feature')\n","Label_path = os.path.join(path, 'Label')\n","os.makedirs(Feature_path, exist_ok=True)\n","os.makedirs(Label_path, exist_ok=True)\n","\n","# 遍历目录中的文件\n","for filename in os.listdir(path):\n","    # 跳过目录，只处理文件\n","    if os.path.isfile(os.path.join(path, filename)):\n","        # 获取文件名前缀\n","        prefix = filename.split('_')[0]\n","        # 检查文件名前缀是否是 area, info, normal 或 centroid\n","        if prefix in ['mesh']:\n","            # 构建源文件和目标文件路径\n","            src_file = os.path.join(path, filename)\n","            dst_file = os.path.join(Feature_path, filename)\n","            # 移动文件\n","            shutil.move(src_file, dst_file)\n","            # print(f'Moved {filename} to {Feature_path}')\n","        elif prefix in ['press']:\n","            # 构建源文件和目标文件路径\n","            src_file = os.path.join(path, filename)\n","            dst_file = os.path.join(Label_path, filename)\n","            # 移动文件\n","            shutil.move(src_file, dst_file)\n","            # print(f'Moved {filename} to {Label_path}')"],"metadata":{"id":"PCb-JOgEWC-u"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import os\n","import shutil\n","\n","# 指定 track_a 目录路径\n","track_a_path = 'Dataset/track_A'\n","\n","# 新建Inference 文件夹\n","\n","inference_path = os.path.join(track_a_path, 'Inference')\n","\n","os.makedirs(inference_path, exist_ok=True)\n","\n","# 遍历 track_b 目录中的文件\n","for filename in os.listdir(track_a_path):\n","    # 跳过目录，只处理文件\n","    if os.path.isfile(os.path.join(track_a_path, filename)):\n","        # 获取文件名前缀\n","        prefix = filename.split('_')[0]\n","        # 检查文件名前缀是否是 area, info, normal 或 centroid\n","        if prefix in ['mesh']:\n","            # 构建源文件和目标文件路径\n","            src_file = os.path.join(track_a_path, filename)\n","            dst_file = os.path.join(inference_path, filename)\n","            # 移动文件\n","            shutil.move(src_file, dst_file)\n","            # print(f'Moved {filename} to {inference_path}')\n"],"metadata":{"id":"coE0Qu-YU8pX"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import os\n","import shutil\n","\n","# 指定 track_b 目录路径\n","track_b_path = 'Dataset/track_B'\n","\n","# 新建 Auxiliary 和 Inference 文件夹\n","auxiliary_path = os.path.join(track_b_path, 'Auxiliary')\n","inference_path = os.path.join(track_b_path, 'Inference')\n","os.makedirs(auxiliary_path, exist_ok=True)\n","os.makedirs(inference_path, exist_ok=True)\n","\n","# 遍历 track_b 目录中的文件\n","for filename in os.listdir(track_b_path):\n","    # 跳过目录，只处理文件\n","    if os.path.isfile(os.path.join(track_b_path, filename)):\n","        # 获取文件名前缀\n","        prefix = filename.split('_')[0]\n","        # 检查文件名前缀是否是 area, info, normal 或 centroid\n","        if prefix in ['area', 'info', 'normal']:\n","            # 构建源文件和目标文件路径\n","            src_file = os.path.join(track_b_path, filename)\n","            dst_file = os.path.join(auxiliary_path, filename)\n","            # 移动文件\n","            shutil.move(src_file, dst_file)\n","            # print(f'Moved {filename} to {auxiliary_path}')\n","        elif prefix in ['centroid','mesh']:\n","            # 构建源文件和目标文件路径\n","            src_file = os.path.join(track_b_path, filename)\n","            dst_file = os.path.join(inference_path, filename)\n","            # 移动文件\n","            shutil.move(src_file, dst_file)\n","            # print(f'Moved {filename} to {inference_path}')\n","\n"],"metadata":{"id":"UhavukX0TFcQ"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# Model"],"metadata":{"id":"VV4alT_R22se"}},{"cell_type":"code","source":["#!/usr/bin/env python\n","# -*- coding:utf-8 _*-\n","import math\n","import numpy as np\n","import torch\n","import torch.nn as nn\n","# import dgl\n","from einops import repeat, rearrange\n","from torch.nn import functional as F\n","from torch.nn import GELU, ReLU, Tanh, Sigmoid\n","from torch.nn.utils.rnn import pad_sequence\n","# from timm.models import trunc_normal_\n","# from utils import MultipleTensors\n","from torch.nn import init\n","import torch.nn as nn\n","import torch.nn.functional as F\n","\n","\n","ACTIVATION = {'gelu':nn.GELU(),'tanh':nn.Tanh(),'sigmoid':nn.Sigmoid(),'relu':nn.ReLU(),'leaky_relu':nn.LeakyReLU(0.1),'softplus':nn.Softplus(),'ELU':nn.ELU()}\n","\n","'''\n","    A simple MLP class, includes at least 2 layers and n hidden layers\n","'''\n","class MLP(nn.Module):\n","    def __init__(self, n_input, n_hidden, n_output, n_layers=1, act='gelu'):\n","        super(MLP, self).__init__()\n","\n","        if act in ACTIVATION.keys():\n","            self.act = ACTIVATION[act]\n","        else:\n","            raise NotImplementedError\n","        self.n_input = n_input\n","        self.n_hidden = n_hidden\n","        self.n_output = n_output\n","        self.n_layers = n_layers\n","        self.linear_pre = nn.Linear(n_input, n_hidden)\n","        self.linear_post = nn.Linear(n_hidden, n_output)\n","        self.linears = nn.ModuleList([nn.Linear(n_hidden, n_hidden) for _ in range(n_layers)])\n","\n","        # self.bns = nn.ModuleList([nn.BatchNorm1d(n_hidden) for _ in range(n_layers)])\n","\n","\n","\n","    def forward(self, x):\n","        x = self.act(self.linear_pre(x))\n","        for i in range(self.n_layers):\n","            x = self.act(self.linears[i](x)) + x\n","            # x = self.act(self.bns[i](self.linears[i](x))) + x\n","\n","        x = self.linear_post(x)\n","        return x\n","\n","\n","class AgentAttention(nn.Module):\n","    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.,\n","                 agent_num=49, **kwargs):\n","        super().__init__()\n","        self.ln = nn.LayerNorm(dim)\n","        self.dim = dim\n","        self.num_heads = num_heads\n","        head_dim = dim // num_heads\n","        self.scale = head_dim ** -0.5\n","        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n","        self.attn_drop = nn.Dropout(attn_drop)\n","        self.proj = nn.Linear(dim, dim)\n","        self.proj_drop = nn.Dropout(proj_drop)\n","        self.softmax = nn.Softmax(dim=-1)\n","        self.agent_num=agent_num\n","        self.agent_token = nn.Parameter(torch.empty(agent_num, dim))\n","\n","    def _initialize_parameters(self):\n","        # 使用 Xavier 正态分布初始化\n","        init.xavier_normal_(self.agent_token)\n","        # 或者使用 Xavier 均匀分布初始化\n","        # init.xavier_uniform_(self.token)\n","\n","    def forward(self, x):\n","        \"\"\"\n","        Args:\n","            x: input features with shape of (num_windows*B, N, C)\n","            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None\n","        \"\"\"\n","        b, n, c = x.shape\n","\n","        num_heads = self.num_heads\n","        head_dim = c // num_heads\n","        qkv = self.qkv(x).reshape(b, n, 3, c).permute(2, 0, 1, 3)\n","        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)\n","        # q, k, v: b, n, c\n","\n","        agent_tokens = self.agent_token.unsqueeze(0).repeat(b,1,1)\n","        agent_tokens=self.ln(agent_tokens)\n","        q = q.reshape(b, n, num_heads, head_dim).permute(0, 2, 1, 3)\n","        k = k.reshape(b, n, num_heads, head_dim).permute(0, 2, 1, 3)\n","        v = v.reshape(b, n, num_heads, head_dim).permute(0, 2, 1, 3)\n","        agent_tokens = agent_tokens.reshape(b, self.agent_num, num_heads, head_dim).permute(0, 2, 1, 3)\n","\n","\n","        agent_attn = self.softmax((agent_tokens * self.scale) @ k.transpose(-2, -1) )\n","        agent_attn = self.attn_drop(agent_attn)\n","        agent_v = agent_attn @ v\n","\n","\n","        q_attn = self.softmax((q * self.scale) @ agent_tokens.transpose(-2, -1) )\n","        q_attn = self.attn_drop(q_attn)\n","        x = q_attn @ agent_v\n","\n","        x = x.transpose(1, 2).reshape(b, n, c)\n","\n","\n","        x = self.proj(x)\n","        x = self.proj_drop(x)\n","        return x\n","\n","class MultipleTensors():\n","    def __init__(self, x):\n","        self.x = x\n","\n","    def to(self, device):\n","        self.x = [x_.to(device) for x_ in self.x]\n","        return self\n","\n","    def __len__(self):\n","        return len(self.x)\n","\n","    def __getitem__(self, item):\n","        return self.x[item]\n","\n","\n","class MoEGPTConfig():\n","    \"\"\" base GPT config, params common to all GPT versions \"\"\"\n","\n","    def __init__(self, attn_type='linear', embd_pdrop=0.0, resid_pdrop=0.0, attn_pdrop=0.0, n_embd=128, n_head=1,\n","                 n_layer=3, block_size=128, n_inner=4, act='gelu', n_experts=2, space_dim=3, branch_sizes=None,\n","                 n_inputs=1, agent_num=64):\n","        self.attn_type = attn_type\n","        self.embd_pdrop = embd_pdrop\n","        self.resid_pdrop = resid_pdrop\n","        self.attn_pdrop = attn_pdrop\n","        self.n_embd = n_embd  # 64\n","        self.n_head = n_head\n","        self.n_layer = n_layer\n","        self.block_size = block_size\n","        self.n_inner = n_inner * self.n_embd\n","        self.act = act\n","        self.n_experts = n_experts\n","        self.space_dim = space_dim\n","        self.branch_sizes = branch_sizes\n","        self.n_inputs = n_inputs\n","        self.agent_num = agent_num\n","\n","\n","\n","\n","\n","class LinearAttention(nn.Module):\n","    \"\"\"\n","    A vanilla multi-head masked self-attention layer with a projection at the end.\n","    It is possible to use torch.nn.MultiheadAttention here but I am including an\n","    explicit implementation here to show that there is nothing too scary here.\n","    \"\"\"\n","\n","    def __init__(self, config):\n","        super(LinearAttention, self).__init__()\n","        assert config.n_embd % config.n_head == 0\n","        # key, query, value projections for all heads\n","        self.key = nn.Linear(config.n_embd, config.n_embd)\n","        self.query = nn.Linear(config.n_embd, config.n_embd)\n","        self.value = nn.Linear(config.n_embd, config.n_embd)\n","        # regularization\n","        self.attn_drop = nn.Dropout(config.attn_pdrop)\n","        # output projection\n","        self.proj = nn.Linear(config.n_embd, config.n_embd)\n","\n","        self.n_head = config.n_head\n","\n","        self.attn_type = 'l1'\n","\n","    '''\n","        Linear Attention and Linear Cross Attention (if y is provided)\n","    '''\n","\n","    def forward(self, x, y=None, layer_past=None):\n","        y = x if y is None else y\n","        B, T1, C = x.size()\n","        _, T2, _ = y.size()\n","        # calculate query, key, values for all heads in batch and move head forward to be the batch dim\n","        q = self.query(x).view(B, T1, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)\n","        k = self.key(y).view(B, T2, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)\n","        v = self.value(y).view(B, T2, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)\n","\n","        if self.attn_type == 'l1':\n","            q = q.softmax(dim=-1)\n","            k = k.softmax(dim=-1)  #\n","            k_cumsum = k.sum(dim=-2, keepdim=True)\n","            D_inv = 1. / (q * k_cumsum).sum(dim=-1, keepdim=True)  # normalized\n","        elif self.attn_type == \"galerkin\":\n","            q = q.softmax(dim=-1)\n","            k = k.softmax(dim=-1)  #\n","            D_inv = 1. / T2  # galerkin\n","        elif self.attn_type == \"l2\":  # still use l1 normalization\n","            q = q / q.norm(dim=-1, keepdim=True, p=1)\n","            k = k / k.norm(dim=-1, keepdim=True, p=1)\n","            k_cumsum = k.sum(dim=-2, keepdim=True)\n","            D_inv = 1. / (q * k_cumsum).abs().sum(dim=-1, keepdim=True)  # 这里有没有做sum都可以，按照论文公式是没有sum  # normalized\n","        else:\n","            raise NotImplementedError\n","\n","        context = k.transpose(-2, -1) @ v\n","        y = self.attn_drop((q @ context) * D_inv + q)\n","\n","        # output projection\n","        y = rearrange(y, 'b h n d -> b n (h d)')\n","        y = self.proj(y)\n","        return y\n","\n","\n","class LinearCrossAttention(nn.Module):\n","    \"\"\"\n","    A vanilla multi-head masked self-attention layer with a projection at the end.\n","    It is possible to use torch.nn.MultiheadAttention here but I am including an\n","    explicit implementation here to show that there is nothing too scary here.\n","    \"\"\"\n","\n","    def __init__(self, config):\n","        super(LinearCrossAttention, self).__init__()\n","        assert config.n_embd % config.n_head == 0\n","        # key, query, value projections for all heads\n","        self.query = nn.Linear(config.n_embd, config.n_embd)\n","        self.keys = nn.ModuleList([nn.Linear(config.n_embd, config.n_embd) for _ in range(config.n_inputs)])\n","        self.values = nn.ModuleList([nn.Linear(config.n_embd, config.n_embd) for _ in range(config.n_inputs)])\n","        # regularization\n","        self.attn_drop = nn.Dropout(config.attn_pdrop)\n","        # output projection\n","        self.proj = nn.Linear(config.n_embd, config.n_embd)\n","\n","        self.n_head = config.n_head\n","        self.n_inputs = config.n_inputs\n","\n","        self.attn_type = 'l1'\n","\n","    '''\n","        Linear Attention and Linear Cross Attention (if y is provided)\n","    '''\n","\n","    def forward(self, x, y=None, layer_past=None):\n","        y = x if y is None else y\n","        B, T1, C = x.size()\n","        # calculate query, key, values for all heads in batch and move head forward to be the batch dim\n","        q = self.query(x).view(B, T1, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)\n","        q = q.softmax(dim=-1)\n","        out = q\n","        for i in range(self.n_inputs):\n","            _, T2, _ = y[i].size()\n","            k = self.keys[i](y[i]).view(B, T2, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)\n","            v = self.values[i](y[i]).view(B, T2, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)\n","            k = k.softmax(dim=-1)  #\n","            k_cumsum = k.sum(dim=-2, keepdim=True)\n","            D_inv = 1. / (q * k_cumsum).sum(dim=-1, keepdim=True)  # normalized\n","            out = out + 1 * (q @ (k.transpose(-2, -1) @ v)) * D_inv\n","\n","        # output projection\n","        out = rearrange(out, 'b h n d -> b n (h d)')\n","        out = self.proj(out)\n","        return out\n","\n","\n","'''\n","    X: N*T*C --> N*(4*n + 3)*C\n","'''\n","\n","\n","def horizontal_fourier_embedding(X, n=3):\n","    freqs = 2 ** torch.linspace(-n, n, 2 * n + 1).to(X.device)\n","    freqs = freqs[None, None, None, ...]\n","    X_ = X.unsqueeze(-1).repeat([1, 1, 1, 2 * n + 1])\n","    X_cos = torch.cos(freqs * X_)\n","    X_sin = torch.sin(freqs * X_)\n","    X = torch.cat([X.unsqueeze(-1), X_cos, X_sin], dim=-1).view(X.shape[0], X.shape[1], -1)\n","    return X\n","\n","\n","'''\n","    Self and Cross Attention block for CGPT, contains  a cross attention block and a self attention block\n","'''\n","\n","\n","class MIOECrossAttentionBlock(nn.Module):\n","    def __init__(self, config):\n","        super(MIOECrossAttentionBlock, self).__init__()\n","        self.ln1 = nn.LayerNorm(config.n_embd)\n","        self.ln2_branch = nn.ModuleList([nn.LayerNorm(config.n_embd) for _ in range(config.n_inputs)])\n","        self.ln3 = nn.LayerNorm(config.n_embd)\n","        self.ln4 = nn.LayerNorm(config.n_embd)\n","        self.ln5 = nn.LayerNorm(config.n_embd)\n","        if config.attn_type == 'linear':\n","            print('Using Linear Attention')\n","            self.selfattn = AgentAttention(config.n_embd,config.n_head,agent_num=config.agent_num)\n","            self.crossattn = LinearCrossAttention(config)\n","        else:\n","            raise NotImplementedError\n","\n","        if config.act == 'gelu':\n","            self.act = GELU\n","        elif config.act == \"tanh\":\n","            self.act = Tanh\n","        elif config.act == 'relu':\n","            self.act = ReLU\n","        elif config.act == 'sigmoid':\n","            self.act = Sigmoid\n","\n","        self.resid_drop1 = nn.Dropout(config.resid_pdrop)\n","        self.resid_drop2 = nn.Dropout(config.resid_pdrop)\n","\n","        self.n_experts = config.n_experts\n","        self.n_inputs = config.n_inputs\n","\n","        self.moe_mlp1 = nn.ModuleList([nn.Sequential(\n","            nn.Linear(config.n_embd, config.n_inner),\n","            self.act(),\n","            nn.Linear(config.n_inner, config.n_embd),\n","        ) for _ in range(self.n_experts)])\n","\n","        self.moe_mlp2 = nn.ModuleList([nn.Sequential(\n","            nn.Linear(config.n_embd, config.n_inner),\n","            self.act(),\n","            nn.Linear(config.n_inner, config.n_embd),\n","        ) for _ in range(self.n_experts)])\n","\n","        self.gatenet = nn.Sequential(\n","            nn.Linear(config.space_dim, config.n_inner),\n","            self.act(),\n","            nn.Linear(config.n_inner, config.n_inner),\n","            self.act(),\n","            nn.Linear(config.n_inner, self.n_experts)\n","        )\n","\n","    def ln_branchs(self, y):\n","        return MultipleTensors([self.ln2_branch[i](y[i]) for i in range(self.n_inputs)])\n","\n","    '''\n","        x: [B, T1, C], y:[B, T2, C], pos:[B, T1, n]\n","    '''\n","\n","    def forward(self, x, y, pos):\n","        gate_score = F.softmax(self.gatenet(pos), dim=-1).unsqueeze(2)  # B, T1, 1, m\n","        x = x + self.resid_drop1(self.crossattn(self.ln1(x), self.ln_branchs(y)))\n","        x_moe1 = torch.stack([self.moe_mlp1[i](x) for i in range(self.n_experts)], dim=-1)  # B, T1, C, m\n","        x_moe1 = (gate_score * x_moe1).sum(dim=-1, keepdim=False)\n","        x = x + self.ln3(x_moe1)\n","        x = x + self.resid_drop2(self.selfattn(self.ln4(x)))\n","        x_moe2 = torch.stack([self.moe_mlp2[i](x) for i in range(self.n_experts)], dim=-1)  # B, T1, C, m\n","        x_moe2 = (gate_score * x_moe2).sum(dim=-1, keepdim=False)\n","        x = x + self.ln5(x_moe2)\n","        return x\n","\n","    #### No layernorm\n","    # def forward(self, x, y):\n","    #     # y = self.selfattn_branch(self.ln5(y))\n","    #     x = x + self.resid_drop1(self.crossattn(x, y))\n","    #     x = x + self.mlp1(x)\n","    #     x = x + self.resid_drop2(self.selfattn(x))\n","    #     x = x + self.mlp2(x)\n","    #\n","    #     return x\n","\n","\n","'''\n","    Cross Attention GPT neural operator\n","    Trunck Net: geom\n","'''\n","\n","\n","class GNOT_Agent(nn.Module):\n","    def __init__(self,\n","                 trunk_size=2,\n","                 branch_sizes=None,  # 分支输入的维度 举例[1，3，3] [sdf,wind,track ...]\n","                 space_dim=2,\n","                 output_size=3,\n","                 n_layers=2,\n","                 n_hidden=64,\n","                 n_head=1,\n","                 n_experts=2,\n","                 n_inner=4,\n","                 mlp_layers=2,\n","                 attn_type='linear',\n","                 act='gelu',\n","                 ffn_dropout=0.0,\n","                 attn_dropout=0.0,\n","                 horiz_fourier_dim=0,\n","                 agent_num=64,\n","                 ):\n","        super(GNOT_Agent, self).__init__()\n","\n","        self.horiz_fourier_dim = horiz_fourier_dim\n","        self.trunk_size = trunk_size * (4 * horiz_fourier_dim + 3) if horiz_fourier_dim > 0 else trunk_size\n","        self.branch_sizes = [bsize * (4 * horiz_fourier_dim + 3) for bsize in\n","                             branch_sizes] if horiz_fourier_dim > 0 else branch_sizes\n","        self.n_inputs = len(self.branch_sizes)\n","        self.output_size = output_size\n","        self.space_dim = space_dim\n","        self.gpt_config = MoEGPTConfig(attn_type=attn_type, embd_pdrop=ffn_dropout, resid_pdrop=ffn_dropout,\n","                                       attn_pdrop=attn_dropout, n_embd=n_hidden, n_head=n_head, n_layer=n_layers,\n","                                       block_size=128, act=act, n_experts=n_experts, space_dim=space_dim,\n","                                       branch_sizes=branch_sizes, n_inputs=len(branch_sizes), n_inner=n_inner,agent_num=agent_num)\n","\n","        self.trunk_mlp = MLP(self.trunk_size, n_hidden, n_hidden, n_layers=mlp_layers, act=act)\n","        self.branch_mlps = nn.ModuleList(\n","            [MLP(bsize, n_hidden, n_hidden, n_layers=mlp_layers, act=act) for bsize in self.branch_sizes])\n","\n","        self.blocks = nn.Sequential(*[MIOECrossAttentionBlock(self.gpt_config) for _ in range(self.gpt_config.n_layer)])\n","\n","        self.out_mlp = MLP(n_hidden, n_hidden, output_size, n_layers=mlp_layers)\n","\n","        # self.apply(self._init_weights)\n","\n","        self.__name__ = 'MIOEGPT'\n","\n","    def _init_weights(self, module):\n","        if isinstance(module, (nn.Linear, nn.Embedding)):\n","            module.weight.data.normal_(mean=0.0, std=0.0002)\n","            if isinstance(module, nn.Linear) and module.bias is not None:\n","                module.bias.data.zero_()\n","        elif isinstance(module, nn.LayerNorm):\n","            module.bias.data.zero_()\n","            module.weight.data.fill_(1.0)\n","\n","    def forward(self, query_points, conditions, deep=None):\n","\n","        pos = query_points[:, :, 0:self.space_dim]\n","\n","        # x = torch.cat([x, u_p.unsqueeze(1).repeat([1, x.shape[1], 1])], dim=-1)\n","\n","        # if self.horiz_fourier_dim > 0:\n","        #     x = horizontal_fourier_embedding(x, self.horiz_fourier_dim)\n","        #     z = horizontal_fourier_embedding(z, self.horiz_fourier_dim)\n","\n","        x = self.trunk_mlp(query_points)\n","        z = MultipleTensors([self.branch_mlps[i](conditions[i]) for i in range(self.n_inputs)])\n","        if deep == None:\n","            for block in self.blocks:\n","                x = block(x, z, pos)\n","            x_out = self.out_mlp(x)\n","        else:\n","            for i in range(deep):\n","                x = self.blocks[i](x, z, pos)\n","            x_out = self.out_mlp(x)\n","\n","        # x_out = torch.cat([x[i, :num] for i, num in enumerate(g.batch_num_nodes())], dim=0)\n","        return x_out\n","\n","\n","\n"],"metadata":{"id":"gS0mLrw-YUeM"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# Args"],"metadata":{"id":"TA4DWJ3Qn60x"}},{"cell_type":"code","source":["#!/usr/bin/env python\n","# -*- coding:utf-8 _*-\n","\n","import argparse\n","\n","\n","def get_args():\n","    parser = argparse.ArgumentParser(description='trackB')\n","    parser.add_argument('--dataset', type=str,\n","                        default='ns2d',\n","                        choices=['heat2d', 'ns2d', 'inductor2d', 'heatsink3d', 'ns2d_time', 'darcy2d', ])\n","    parser.add_argument('--train_dir', default='./Dataset/data', type=str)\n","    parser.add_argument('--test_dir', default='./Dataset/track_A', type=str)\n","    parser.add_argument('--wind_direction', default='z', type=str)\n","    parser.add_argument('--track', default='track_A', type=str)\n","    parser.add_argument('--train_log', default='train', type=str)\n","    parser.add_argument('--expand_data', type=int,default=0)\n","    parser.add_argument('--val_ratio', default=0.2, type=float)\n","    parser.add_argument('--p', default=0.5, type=float,help='p of using tracka')\n","\n","    parser.add_argument('--iter_per_epoch', default=150, type=int)\n","    parser.add_argument('--cat_area', type=int, default=1,\n","                        help='cat area or use condition')\n","    parser.add_argument('--use_transform', type=int, default=0)\n","    parser.add_argument('--component', type=str,\n","                        default='all', )\n","\n","    parser.add_argument('--seed', type=int, default=2024, metavar='Seed',\n","                        help='random seed (default: 1127802)')\n","\n","    parser.add_argument('--gpu', type=int, default=0, help='gpu id')\n","    parser.add_argument('--use-tb', type=int, default=1, help='whether use tensorboard')\n","    parser.add_argument('--comment', type=str, default=\"\", help=\"comment for the experiment\")\n","\n","    parser.add_argument('--train-num', type=str, default='all')\n","    parser.add_argument('--test-num', type=str, default='all')\n","\n","    parser.add_argument('--sort-data', type=int, default=0)\n","\n","    parser.add_argument('--normalize_x', type=str, default='unit',\n","                        choices=['none', 'minmax', 'unit'])\n","    parser.add_argument('--use-normalizer', type=str, default='unit',\n","                        choices=['none', 'minmax', 'unit', 'quantile'],\n","                        help=\"whether normalize y\")\n","\n","    parser.add_argument('--epochs', type=int, default=200, metavar='N',\n","                        help='number of epochs to train (default: 100)')\n","    parser.add_argument('--val_epochs', type=int, default=10,)\n","    parser.add_argument('--optimizer', type=str, default='AdamW', choices=['Adam', 'AdamW'])\n","\n","    parser.add_argument('--lr', type=float, default=0.001, metavar='LR',\n","                        help='max learning rate (default: 0.001)')\n","    parser.add_argument('--weight-decay', type=float, default=5e-5\n","                        )\n","    parser.add_argument('--grad-clip', type=str, default=1000.0\n","                        )\n","    parser.add_argument('--batch-size', type=int, default=1, metavar='bsz',\n","                        help='input batch size for training (default: 8)')\n","    parser.add_argument('--val-batch-size', type=int, default=8, metavar='bsz',\n","                        help='input batch size for validation (default: 4)')\n","    parser.add_argument('--no-cuda', default=0,type=int,\n","                        help='disables CUDA training')\n","\n","\n","    parser.add_argument('--lr-method', type=str, default='cycle',\n","                        choices=['cycle', 'step', 'warmup'])\n","    parser.add_argument('--lr-step-size', type=int, default=20\n","                        )\n","    parser.add_argument('--warmup-epochs', type=int, default=50)\n","\n","    parser.add_argument('--loss-name', type=str, default='rel2',\n","                        choices=['rel2', 'rel1', 'l2', 'l1'])\n","    #### public model architecture parameters\n","\n","    parser.add_argument('--model-name', type=str, default='GNOT',\n","                        choices=['CGPT', 'GNOT', ])\n","    parser.add_argument('--n-hidden', type=int, default=32)\n","    parser.add_argument('--n-layers', type=int, default=5)\n","\n","    #### MLP parameters\n","\n","    # common\n","    parser.add_argument('--act', type=str, default='gelu', choices=['gelu', 'relu', 'tanh', 'sigmoid'])\n","    parser.add_argument('--loss', type=str, default='default', choices=['default','mse', 'l1', 'huber', 'smoothl1'])\n","    parser.add_argument('--model', type=str, default='GAGENT')\n","    parser.add_argument('--n-head', type=int, default=2)\n","    parser.add_argument('--slice_num', type=int, default=64)\n","    parser.add_argument('--agent_num', type=int, default=32)\n","    parser.add_argument('--ffn-dropout', type=float, default=0.0, metavar='ffn_dropout',\n","                        help='dropout for the FFN in attention (default: 0.0)')\n","    parser.add_argument('--attn-dropout', type=float, default=0.0)\n","    parser.add_argument('--mlp-layers', type=int, default=3)\n","\n","    # Transformer\n","    # parser.add_argument('--subsampled-len',type=int, default=256)\n","    parser.add_argument('--attn-type', type=str, default='linear',\n","                        choices=['random', 'linear', 'gated', 'hydra', 'kernel'])\n","    parser.add_argument('--hfourier-dim', type=int, default=0)\n","\n","    # GNOT\n","    parser.add_argument('--n_experts', type=int, default=2)\n","    parser.add_argument('--input_dim', type=int, default=4)\n","    parser.add_argument('--output_dim', type=int, default=1)\n","    parser.add_argument('--space_dim', type=int, default=3)\n","\n","    parser.add_argument('--branch_sizes', nargs=\"*\", type=int, default=[3, 3])\n","    parser.add_argument('--n-inner', type=int, default=4)\n","\n","    return parser.parse_known_args()[0]\n"],"metadata":{"id":"KmeREhHyjUJP"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# Utils"],"metadata":{"id":"UQxAwmUYnWkp"}},{"cell_type":"code","source":["#!/usr/bin/env python\n","# -*- coding:utf-8 _*-\n","import os\n","import torch\n","import numpy as np\n","import operator\n","import matplotlib.pyplot as plt\n","import torch\n","import numpy as np\n","import torch.special as ts\n","from torch import nn\n","\n","from scipy import interpolate\n","from functools import reduce\n","from pathlib import Path\n","from torch.nn import  MSELoss,L1Loss,HuberLoss,SmoothL1Loss\n","class LpLoss(nn.Module):\n","    def __init__(self, d=2, p=2, size_average=True, reduction='mean'):\n","        super(LpLoss, self).__init__()\n","        assert d > 0 and p > 0\n","        self.d = d\n","        self.p = p\n","        self.reduction = reduction\n","        self.size_average = size_average\n","\n","    def forward(self, x, y):\n","        num_examples = x.size(0)\n","        h = 1.0 / (x.size(1) - 1.0)\n","        all_norms = (h ** (self.d / self.p)) * torch.norm(\n","            x.reshape(num_examples, -1) - y.reshape(num_examples, -1), self.p, dim=1\n","        )\n","        if self.reduction == 'mean':\n","            return all_norms.mean() if self.size_average else all_norms.sum()\n","        elif self.reduction == 'sum':\n","            return all_norms.sum()\n","        else:\n","            return all_norms\n","\n","def extract_condition(data):\n","\n","    keys_to_keep = {'vertices', 'pressure','area'}\n","\n","    filtered_data = [value for key, value in data.items() if key not in keys_to_keep]\n","\n","    return filtered_data\n","\n","\n","\n","\n","def get_model(args):\n","    trunk_size,  branch_sizes, output_size = args.input_dim, args.branch_sizes, args.output_dim\n","    output_size = args.output_dim\n","    agent_num=args.agent_num\n","    if args.model=='GNOT':\n","      print(1)\n","        # return GNOT(trunk_size=trunk_size, branch_sizes=branch_sizes, output_size=output_size, n_layers=args.n_layers,\n","        #             n_hidden=args.n_hidden, n_head=args.n_head, attn_type=args.attn_type, ffn_dropout=args.ffn_dropout,\n","        #             attn_dropout=args.attn_dropout, mlp_layers=args.mlp_layers, act=args.act,\n","        #             horiz_fourier_dim=args.hfourier_dim, space_dim=args.space_dim, n_experts=args.n_experts,\n","        #             n_inner=args.n_inner)\n","    else:\n","        return GNOT_Agent(trunk_size=trunk_size, branch_sizes=branch_sizes, output_size=output_size, n_layers=args.n_layers,\n","                    n_hidden=args.n_hidden, n_head=args.n_head, attn_type=args.attn_type, ffn_dropout=args.ffn_dropout,\n","                    attn_dropout=args.attn_dropout, mlp_layers=args.mlp_layers, act=args.act,\n","                    horiz_fourier_dim=args.hfourier_dim, space_dim=args.space_dim, n_experts=args.n_experts,\n","                    n_inner=args.n_inner,agent_num=agent_num)\n","\n","def get_dataset(args):\n","    return CarDataset(\n","        Path(args.train_dir),\n","        Path(args.test_dir),\n","        expand_data=args.expand_data,\n","        track=args.track,\n","        wind_direction=args.wind_direction,\n","        val_ratio=args.val_ratio,\n","        use_transform=args.use_transform\n","\n","    )\n","\n","def get_loss(args):\n","    if args.loss=='default':\n","        return LpLoss()\n","    if args.loss=='mse':\n","        return MSELoss()\n","    if args.loss=='l1':\n","        return L1Loss()\n","    if args.loss=='huber':\n","        return HuberLoss()\n","    if args.loss=='smoothl1':\n","        return SmoothL1Loss()\n","\n","def get_seed(s, printout=True, cudnn=True):\n","    # rd.seed(s)\n","    os.environ['PYTHONHASHSEED'] = str(s)\n","    np.random.seed(s)\n","    # pd.core.common.random_state(s)\n","    # Torch\n","    torch.manual_seed(s)\n","    torch.cuda.manual_seed(s)\n","    if cudnn:\n","        torch.backends.cudnn.deterministic = True\n","        torch.backends.cudnn.benchmark = False\n","    if torch.cuda.is_available():\n","        torch.cuda.manual_seed_all(s)\n","\n","    message = f'''\n","    os.environ['PYTHONHASHSEED'] = str({s})\n","    numpy.random.seed({s})\n","    torch.manual_seed({s})\n","    torch.cuda.manual_seed({s})\n","    torch.backends.cudnn.deterministic = True\n","    torch.backends.cudnn.benchmark = False\n","    if torch.cuda.is_available():\n","        torch.cuda.manual_seed_all({s})\n","    '''\n","    if printout:\n","        print(\"\\n\")\n","        print(f\"The following code snippets have been run.\")\n","        print(\"=\" * 50)\n","        print(message)\n","        print(\"=\" * 50)\n","\n","\n","def get_num_params(model):\n","    '''\n","    a single entry in cfloat and cdouble count as two parameters\n","    see https://github.com/pytorch/pytorch/issues/57518\n","    '''\n","    # model_parameters = filter(lambda p: p.requires_grad, model.parameters())\n","    # num_params = 0\n","    # for p in model_parameters:\n","    #     # num_params += np.prod(p.size()+(2,) if p.is_complex() else p.size())\n","    #     num_params += p.numel() * (1 + p.is_complex())\n","    # return num_params\n","\n","    c = 0\n","    for p in list(model.parameters()):\n","        c += reduce(operator.mul, list(p.size() + (2,) if p.is_complex() else p.size()))  #### there is complex weight\n","    return c\n","\n","\n","### x: list of tensors\n","class MultipleTensors():\n","    def __init__(self, x):\n","        self.x = x\n","\n","    def to(self, device):\n","        self.x = [x_.to(device) for x_ in self.x]\n","        return self\n","\n","    def __len__(self):\n","        return len(self.x)\n","\n","    def __getitem__(self, item):\n","        return self.x[item]\n","\n","\n","# whether need to transpose\n","def plot_heatmap(\n","        x, y, z, path=None, vmin=None, vmax=None, cmap=None,\n","        title=\"\", xlabel=\"x\", ylabel=\"y\", show=False\n","):\n","    '''\n","    Plot heat map for a 3-dimension data\n","    '''\n","    plt.cla()\n","    # plt.figure()\n","    xx = np.linspace(np.min(x), np.max(x))\n","    yy = np.linspace(np.min(y), np.max(y))\n","    xx, yy = np.meshgrid(xx, yy)\n","\n","    vals = interpolate.griddata(np.array([x, y]).T, np.array(z),\n","                                (xx, yy), method='cubic')\n","    vals_0 = interpolate.griddata(np.array([x, y]).T, np.array(z),\n","                                  (xx, yy), method='nearest')\n","    vals[np.isnan(vals)] = vals_0[np.isnan(vals)]\n","\n","    if vmin is not None and vmax is not None:\n","        fig = plt.imshow(vals,\n","                         extent=[np.min(x), np.max(x), np.min(y), np.max(y)],\n","                         aspect=\"equal\", interpolation=\"bicubic\", cmap=cmap,\n","                         vmin=vmin, vmax=vmax, origin='lower')\n","    elif vmin is not None:\n","        fig = plt.imshow(vals,\n","                         extent=[np.min(x), np.max(x), np.min(y), np.max(y)],\n","                         aspect=\"equal\", interpolation=\"bicubic\", cmap=cmap,\n","                         vmin=vmin, origin='lower')\n","    else:\n","        fig = plt.imshow(vals,\n","                         extent=[np.min(x), np.max(x), np.min(y), np.max(y)], cmap=cmap,\n","                         aspect=\"equal\", interpolation=\"bicubic\", origin='lower')\n","    fig.axes.set_autoscale_on(False)\n","    plt.xlabel(xlabel)\n","    plt.ylabel(ylabel)\n","    plt.title(title)\n","    plt.colorbar()\n","    if path:\n","        plt.savefig(path)\n","    if show:\n","        plt.show()\n","    plt.close()\n","\n","\n","import contextlib\n","\n","\n","class Interp1d(torch.autograd.Function):\n","    def __call__(self, x, y, xnew, out=None):\n","        return self.forward(x, y, xnew, out)\n","\n","    def forward(ctx, x, y, xnew, out=None):\n","        \"\"\"\n","        Linear 1D interpolation on the GPU for Pytorch.\n","        This function returns interpolated values of a set of 1-D functions at\n","        the desired query points `xnew`.\n","        This function is working similarly to Matlab™ or scipy functions with\n","        the `linear` interpolation mode on, except that it parallelises over\n","        any number of desired interpolation problems.\n","        The code will run on GPU if all the tensors provided are on a cuda\n","        device.\n","        Parameters\n","        ----------\n","        x : (N, ) or (D, N) Pytorch Tensor\n","            A 1-D or 2-D tensor of real values.\n","        y : (N,) or (D, N) Pytorch Tensor\n","            A 1-D or 2-D tensor of real values. The length of `y` along its\n","            last dimension must be the same as that of `x`\n","        xnew : (P,) or (D, P) Pytorch Tensor\n","            A 1-D or 2-D tensor of real values. `xnew` can only be 1-D if\n","            _both_ `x` and `y` are 1-D. Otherwise, its length along the first\n","            dimension must be the same as that of whichever `x` and `y` is 2-D.\n","        out : Pytorch Tensor, same shape as `xnew`\n","            Tensor for the output. If None: allocated automatically.\n","        \"\"\"\n","        # making the vectors at least 2D\n","        is_flat = {}\n","        require_grad = {}\n","        v = {}\n","        device = []\n","        eps = torch.finfo(y.dtype).eps\n","        for name, vec in {'x': x, 'y': y, 'xnew': xnew}.items():\n","            assert len(vec.shape) <= 2, 'interp1d: all inputs must be ' \\\n","                                        'at most 2-D.'\n","            if len(vec.shape) == 1:\n","                v[name] = vec[None, :]\n","            else:\n","                v[name] = vec\n","            is_flat[name] = v[name].shape[0] == 1\n","            require_grad[name] = vec.requires_grad\n","            device = list(set(device + [str(vec.device)]))\n","        assert len(device) == 1, 'All parameters must be on the same device.'\n","        device = device[0]\n","\n","        # Checking for the dimensions\n","        assert (v['x'].shape[1] == v['y'].shape[1]\n","                and (\n","                        v['x'].shape[0] == v['y'].shape[0]\n","                        or v['x'].shape[0] == 1\n","                        or v['y'].shape[0] == 1\n","                )\n","                ), (\"x and y must have the same number of columns, and either \"\n","                    \"the same number of row or one of them having only one \"\n","                    \"row.\")\n","\n","        reshaped_xnew = False\n","        if ((v['x'].shape[0] == 1) and (v['y'].shape[0] == 1)\n","                and (v['xnew'].shape[0] > 1)):\n","            # if there is only one row for both x and y, there is no need to\n","            # loop over the rows of xnew because they will all have to face the\n","            # same interpolation problem. We should just stack them together to\n","            # call interp1d and put them back in place afterwards.\n","            original_xnew_shape = v['xnew'].shape\n","            v['xnew'] = v['xnew'].contiguous().view(1, -1)\n","            reshaped_xnew = True\n","\n","        # identify the dimensions of output and check if the one provided is ok\n","        D = max(v['x'].shape[0], v['xnew'].shape[0])\n","        shape_ynew = (D, v['xnew'].shape[-1])\n","        if out is not None:\n","            if out.numel() != shape_ynew[0] * shape_ynew[1]:\n","                # The output provided is of incorrect shape.\n","                # Going for a new one\n","                out = None\n","            else:\n","                ynew = out.reshape(shape_ynew)\n","        if out is None:\n","            ynew = torch.zeros(*shape_ynew, device=device)\n","\n","        # moving everything to the desired device in case it was not there\n","        # already (not handling the case things do not fit entirely, user will\n","        # do it if required.)\n","        for name in v:\n","            v[name] = v[name].to(device)\n","\n","        # calling searchsorted on the x values.\n","        ind = ynew.long()\n","\n","        # expanding xnew to match the number of rows of x in case only one xnew is\n","        # provided\n","        if v['xnew'].shape[0] == 1:\n","            v['xnew'] = v['xnew'].expand(v['x'].shape[0], -1)\n","\n","        torch.searchsorted(v['x'].contiguous(),\n","                           v['xnew'].contiguous(), out=ind)\n","\n","        # the `-1` is because searchsorted looks for the index where the values\n","        # must be inserted to preserve order. And we want the index of the\n","        # preceeding value.\n","        ind -= 1\n","        # we clamp the index, because the number of intervals is x.shape-1,\n","        # and the left neighbour should hence be at most number of intervals\n","        # -1, i.e. number of columns in x -2\n","        ind = torch.clamp(ind, 0, v['x'].shape[1] - 1 - 1)\n","\n","        # helper function to select stuff according to the found indices.\n","        def sel(name):\n","            if is_flat[name]:\n","                return v[name].contiguous().view(-1)[ind]\n","            return torch.gather(v[name], 1, ind)\n","\n","        # activating gradient storing for everything now\n","        enable_grad = False\n","        saved_inputs = []\n","        for name in ['x', 'y', 'xnew']:\n","            if require_grad[name]:\n","                enable_grad = True\n","                saved_inputs += [v[name]]\n","            else:\n","                saved_inputs += [None, ]\n","        # assuming x are sorted in the dimension 1, computing the slopes for\n","        # the segments\n","        is_flat['slopes'] = is_flat['x']\n","        # now we have found the indices of the neighbors, we start building the\n","        # output. Hence, we start also activating gradient tracking\n","        with torch.enable_grad() if enable_grad else contextlib.suppress():\n","            v['slopes'] = (\n","                    (v['y'][:, 1:] - v['y'][:, :-1])\n","                    /\n","                    (eps + (v['x'][:, 1:] - v['x'][:, :-1]))\n","            )\n","\n","            # now build the linear interpolation\n","            ynew = sel('y') + sel('slopes') * (\n","                    v['xnew'] - sel('x'))\n","\n","            if reshaped_xnew:\n","                ynew = ynew.view(original_xnew_shape)\n","\n","        ctx.save_for_backward(ynew, *saved_inputs)\n","        return ynew\n","\n","    @staticmethod\n","    def backward(ctx, grad_out):\n","        inputs = ctx.saved_tensors[1:]\n","        gradients = torch.autograd.grad(\n","            ctx.saved_tensors[0],\n","            [i for i in inputs if i is not None],\n","            grad_out, retain_graph=True)\n","        result = [None, ] * 5\n","        pos = 0\n","        for index in range(len(inputs)):\n","            if inputs[index] is not None:\n","                result[index] = gradients[pos]\n","                pos += 1\n","        return (*result,)\n","\n","\n","class TorchQuantileTransformer():\n","    '''\n","    QuantileTransformer implemented by PyTorch\n","    '''\n","\n","    def __init__(\n","            self,\n","            output_distribution,\n","            references_,\n","            quantiles_,\n","            device=torch.device('cpu')\n","    ) -> None:\n","        self.quantiles_ = torch.Tensor(quantiles_).to(device)\n","        self.output_distribution = output_distribution\n","        self._norm_pdf_C = np.sqrt(2 * np.pi)\n","        self.references_ = torch.Tensor(references_).to(device)\n","        BOUNDS_THRESHOLD = 1e-7\n","        self.clip_min = self.norm_ppf(torch.Tensor([BOUNDS_THRESHOLD - np.spacing(1)]))\n","        self.clip_max = self.norm_ppf(torch.Tensor([1 - (BOUNDS_THRESHOLD - np.spacing(1))]))\n","\n","    def norm_pdf(self, x):\n","        return torch.exp(-x ** 2 / 2.0) / self._norm_pdf_C\n","\n","    @staticmethod\n","    def norm_cdf(x):\n","        return ts.ndtr(x)\n","\n","    @staticmethod\n","    def norm_ppf(x):\n","        return ts.ndtri(x)\n","\n","    def transform_col(self, X_col, quantiles, inverse):\n","        BOUNDS_THRESHOLD = 1e-7\n","        output_distribution = self.output_distribution\n","\n","        if not inverse:\n","            lower_bound_x = quantiles[0]\n","            upper_bound_x = quantiles[-1]\n","            lower_bound_y = 0\n","            upper_bound_y = 1\n","        else:\n","            lower_bound_x = 0\n","            upper_bound_x = 1\n","            lower_bound_y = quantiles[0]\n","            upper_bound_y = quantiles[-1]\n","            # for inverse transform, match a uniform distribution\n","            with np.errstate(invalid=\"ignore\"):  # hide NaN comparison warnings\n","                if output_distribution == \"normal\":\n","                    X_col = self.norm_cdf(X_col)\n","                # else output distribution is already a uniform distribution\n","\n","        # find index for lower and higher bounds\n","        with np.errstate(invalid=\"ignore\"):  # hide NaN comparison warnings\n","            if output_distribution == \"normal\":\n","                lower_bounds_idx = X_col - BOUNDS_THRESHOLD < lower_bound_x\n","                upper_bounds_idx = X_col + BOUNDS_THRESHOLD > upper_bound_x\n","            if output_distribution == \"uniform\":\n","                lower_bounds_idx = X_col == lower_bound_x\n","                upper_bounds_idx = X_col == upper_bound_x\n","\n","        isfinite_mask = ~torch.isnan(X_col)\n","        X_col_finite = X_col[isfinite_mask]\n","        torch_interp = Interp1d()\n","        X_col_out = X_col.clone()\n","        if not inverse:\n","            # Interpolate in one direction and in the other and take the\n","            # mean. This is in case of repeated values in the features\n","            # and hence repeated quantiles\n","            #\n","            # If we don't do this, only one extreme of the duplicated is\n","            # used (the upper when we do ascending, and the\n","            # lower for descending). We take the mean of these two\n","            X_col_out[isfinite_mask] = 0.5 * (\n","                    torch_interp(quantiles, self.references_, X_col_finite)\n","                    - torch_interp(-torch.flip(quantiles, [0]), -torch.flip(self.references_, [0]), -X_col_finite)\n","            )\n","        else:\n","            X_col_out[isfinite_mask] = torch_interp(self.references_, quantiles, X_col_finite)\n","\n","        X_col_out[upper_bounds_idx] = upper_bound_y\n","        X_col_out[lower_bounds_idx] = lower_bound_y\n","        # for forward transform, match the output distribution\n","        if not inverse:\n","            with np.errstate(invalid=\"ignore\"):  # hide NaN comparison warnings\n","                if output_distribution == \"normal\":\n","                    X_col_out = self.norm_ppf(X_col_out)\n","                    # find the value to clip the data to avoid mapping to\n","                    # infinity. Clip such that the inverse transform will be\n","                    # consistent\n","                    X_col_out = torch.clip(X_col_out, self.clip_min, self.clip_max)\n","                # else output distribution is uniform and the ppf is the\n","                # identity function so we let X_col unchanged\n","\n","        return X_col_out\n","\n","    def transform(self, X, inverse=True, component='all'):\n","        X_out = torch.zeros_like(X, requires_grad=False)\n","        for feature_idx in range(X.shape[1]):\n","            X_out[:, feature_idx] = self.transform_col(\n","                X[:, feature_idx], self.quantiles_[:, feature_idx], inverse\n","            )\n","        return X_out\n","\n","    def to(self, device):\n","        self.quantiles_ = self.quantiles_.to(device)\n","        self.references_ = self.references_.to(device)\n","        return self\n","\n","\n","'''\n","    Simple normalization layer\n","'''\n","\n","\n","class UnitTransformer():\n","    def __init__(self, X):\n","        self.mean = X.mean(dim=0, keepdim=True)\n","        self.std = X.std(dim=0, keepdim=True) + 1e-8\n","\n","    def to(self, device):\n","        self.mean = self.mean.to(device)\n","        self.std = self.std.to(device)\n","        return self\n","\n","    def transform(self, X, inverse=True, component='all'):\n","        if component == 'all' or 'all-reduce':\n","            if inverse:\n","                orig_shape = X.shape\n","                return (X * (self.std - 1e-8) + self.mean).view(orig_shape)\n","            else:\n","                return (X - self.mean) / self.std\n","        else:\n","            if inverse:\n","                orig_shape = X.shape\n","                return (X * (self.std[:, component] - 1e-8) + self.mean[:, component]).view(orig_shape)\n","            else:\n","                return (X - self.mean[:, component]) / self.std[:, component]\n","\n","\n","'''\n","    Simple pointwise normalization layer, all data must contain the same length, used only for FNO datasets\n","    X: B, N, C\n","'''\n","\n","\n","class PointWiseUnitTransformer():\n","    def __init__(self, X):\n","        self.mean = X.mean(dim=0, keepdim=False)\n","        self.std = X.std(dim=0, keepdim=False) + 1e-8\n","\n","    def to(self, device):\n","        self.mean = self.mean.to(device)\n","        self.std = self.std.to(device)\n","        return self\n","\n","    def transform(self, X, inverse=True, component='all'):\n","        if component == 'all' or 'all-reduce':\n","            if inverse:\n","                orig_shape = X.shape\n","                X = X.view(-1, self.mean.shape[0], self.mean.shape[1])  ### align shape for flat tensor\n","                return (X * (self.std - 1e-8) + self.mean).view(orig_shape)\n","            else:\n","                return (X - self.mean) / self.std\n","        else:\n","            if inverse:\n","                orig_shape = X.shape\n","                X = X.view(-1, self.mean.shape[0], self.mean.shape[1])\n","                return (X * (self.std[:, component] - 1e-8) + self.mean[:, component]).view(orig_shape)\n","            else:\n","                return (X - self.mean[:, component]) / self.std[:, component]\n","\n","\n","'''\n","    x: B, N (not necessary sorted)\n","    y: B, N, C (not necessary sorted)\n","    xnew: B, N (sorted)\n","'''\n","\n","\n","def binterp1d(x, y, xnew, eps=1e-9):\n","    x_, x_indice = torch.sort(x, dim=-1)\n","    y_ = y[torch.arange(x_.shape[0]).unsqueeze(1), x_indice]\n","\n","    x_, y_, xnew = x_.contiguous(), y_.contiguous(), xnew.contiguous()\n","\n","    ind = torch.searchsorted(x_, xnew)\n","    ind -= 1\n","    ind = torch.clamp(ind, 0, x_.shape[1] - 1 - 1)\n","    ind = ind.unsqueeze(-1).repeat([1, 1, y_.shape[-1]])\n","    x_ = x_.unsqueeze(-1).repeat([1, 1, y_.shape[-1]])\n","\n","    slopes = ((y_[:, 1:] - y_[:, :-1]) / (eps + (x_[:, 1:] - x_[:, :-1])))\n","\n","    y_sel = torch.gather(y_, 1, ind)\n","    x_sel = torch.gather(x_, 1, ind)\n","    slopes_sel = torch.gather(slopes, 1, ind)\n","\n","    ynew = y_sel + slopes_sel * (xnew.unsqueeze(-1) - x_sel)\n","\n","    return ynew\n"],"metadata":{"id":"hFW1jQ1VnUt0"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import math\n","import torch\n","\n","from torch import Tensor\n","from typing import List, Optional\n","from torch.optim.optimizer import Optimizer\n","\n","\n","def adam(params: List[Tensor],\n","         grads: List[Tensor],\n","         exp_avgs: List[Tensor],\n","         exp_avg_sqs: List[Tensor],\n","         max_exp_avg_sqs: List[Tensor],\n","         state_steps: List[int],\n","         *,\n","         amsgrad: bool,\n","         beta1: float,\n","         beta2: float,\n","         lr: float,\n","         weight_decay: float,\n","         eps: float):\n","    r\"\"\"Functional API that performs Adam algorithm computation.\n","    See :class:`~torch.optim.Adam` for details.\n","    \"\"\"\n","\n","    for i, param in enumerate(params):\n","\n","        grad = grads[i]\n","        exp_avg = exp_avgs[i]\n","        exp_avg_sq = exp_avg_sqs[i]\n","        step = state_steps[i]\n","\n","        bias_correction1 = 1 - beta1 ** step\n","        bias_correction2 = 1 - beta2 ** step\n","\n","        if weight_decay != 0:\n","            grad = grad.add(param, alpha=weight_decay)\n","\n","        # Decay the first and second moment running average coefficient\n","        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)\n","        exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)\n","        if amsgrad:\n","            # Maintains the maximum of all 2nd moment running avg. till now\n","            torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i])\n","            # Use the max. for normalizing running avg. of gradient\n","            denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps)\n","        else:\n","            denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)\n","\n","        step_size = lr / bias_correction1\n","\n","        param.addcdiv_(exp_avg, denom, value=-step_size)\n","\n","\n","class Adam(Optimizer):\n","    r\"\"\"Implements Adam algorithm.\n","    It has been proposed in `Adam: A Method for Stochastic Optimization`_.\n","    The implementation of the L2 penalty follows changes proposed in\n","    `Decoupled Weight Decay Regularization`_.\n","    Args:\n","        params (iterable): iterable of parameters to optimize or dicts defining\n","            parameter groups\n","        lr (float, optional): learning rate (default: 1e-3)\n","        betas (Tuple[float, float], optional): coefficients used for computing\n","            running averages of gradient and its square (default: (0.9, 0.999))\n","        eps (float, optional): term added to the denominator to improve\n","            numerical stability (default: 1e-8)\n","        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)\n","        amsgrad (boolean, optional): whether to use the AMSGrad variant of this\n","            algorithm from the paper `On the Convergence of Adam and Beyond`_\n","            (default: False)\n","    .. _Adam\\: A Method for Stochastic Optimization:\n","        https://arxiv.org/abs/1412.6980\n","    .. _Decoupled Weight Decay Regularization:\n","        https://arxiv.org/abs/1711.05101\n","    .. _On the Convergence of Adam and Beyond:\n","        https://openreview.net/forum?id=ryQu7f-RZ\n","    \"\"\"\n","\n","    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,\n","                 weight_decay=0, amsgrad=False):\n","        if not 0.0 <= lr:\n","            raise ValueError(\"Invalid learning rate: {}\".format(lr))\n","        if not 0.0 <= eps:\n","            raise ValueError(\"Invalid epsilon value: {}\".format(eps))\n","        if not 0.0 <= betas[0] < 1.0:\n","            raise ValueError(\"Invalid beta parameter at index 0: {}\".format(betas[0]))\n","        if not 0.0 <= betas[1] < 1.0:\n","            raise ValueError(\"Invalid beta parameter at index 1: {}\".format(betas[1]))\n","        if not 0.0 <= weight_decay:\n","            raise ValueError(\"Invalid weight_decay value: {}\".format(weight_decay))\n","        defaults = dict(lr=lr, betas=betas, eps=eps,\n","                        weight_decay=weight_decay, amsgrad=amsgrad)\n","        super(Adam, self).__init__(params, defaults)\n","\n","    def __setstate__(self, state):\n","        super(Adam, self).__setstate__(state)\n","        for group in self.param_groups:\n","            group.setdefault('amsgrad', False)\n","\n","    @torch.no_grad()\n","    def step(self, closure=None):\n","        \"\"\"Performs a single optimization step.\n","        Args:\n","            closure (callable, optional): A closure that reevaluates the model\n","                and returns the loss.\n","        \"\"\"\n","        loss = None\n","        if closure is not None:\n","            with torch.enable_grad():\n","                loss = closure()\n","\n","        for group in self.param_groups:\n","            params_with_grad = []\n","            grads = []\n","            exp_avgs = []\n","            exp_avg_sqs = []\n","            max_exp_avg_sqs = []\n","            state_steps = []\n","            beta1, beta2 = group['betas']\n","\n","            for p in group['params']:\n","                if p.grad is not None:\n","                    params_with_grad.append(p)\n","                    if p.grad.is_sparse:\n","                        raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')\n","                    grads.append(p.grad)\n","\n","                    state = self.state[p]\n","                    # Lazy state initialization\n","                    if len(state) == 0:\n","                        state['step'] = 0\n","                        # Exponential moving average of gradient values\n","                        state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n","                        # Exponential moving average of squared gradient values\n","                        state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n","                        if group['amsgrad']:\n","                            # Maintains max of all exp. moving avg. of sq. grad. values\n","                            state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n","\n","                    exp_avgs.append(state['exp_avg'])\n","                    exp_avg_sqs.append(state['exp_avg_sq'])\n","\n","                    if group['amsgrad']:\n","                        max_exp_avg_sqs.append(state['max_exp_avg_sq'])\n","\n","                    # update the steps for each param group update\n","                    state['step'] += 1\n","                    # record the step after step update\n","                    state_steps.append(state['step'])\n","\n","            adam(params_with_grad,\n","                 grads,\n","                 exp_avgs,\n","                 exp_avg_sqs,\n","                 max_exp_avg_sqs,\n","                 state_steps,\n","                 amsgrad=group['amsgrad'],\n","                 beta1=beta1,\n","                 beta2=beta2,\n","                 lr=group['lr'],\n","                 weight_decay=group['weight_decay'],\n","                 eps=group['eps'])\n","        return loss\n","\n","\n","\n","\n","\n","def adamw(params: List[Tensor],\n","          grads: List[Tensor],\n","          exp_avgs: List[Tensor],\n","          exp_avg_sqs: List[Tensor],\n","          max_exp_avg_sqs: List[Tensor],\n","          state_steps: List[int],\n","          *,\n","          amsgrad: bool,\n","          beta1: float,\n","          beta2: float,\n","          lr: float,\n","          weight_decay: float,\n","          eps: float):\n","    r\"\"\"Functional API that performs AdamW algorithm computation.\n","\n","    See :class:`~torch.optim.AdamW` for details.\n","    \"\"\"\n","    for i, param in enumerate(params):\n","        grad = grads[i]\n","        exp_avg = exp_avgs[i]\n","        exp_avg_sq = exp_avg_sqs[i]\n","        step = state_steps[i]\n","\n","        # Perform stepweight decay\n","        param.mul_(1 - lr * weight_decay)\n","\n","        bias_correction1 = 1 - beta1 ** step\n","        bias_correction2 = 1 - beta2 ** step\n","\n","        # Decay the first and second moment running average coefficient\n","        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)\n","        exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)\n","        if amsgrad:\n","            # Maintains the maximum of all 2nd moment running avg. till now\n","            torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i])\n","            # Use the max. for normalizing running avg. of gradient\n","            denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps)\n","        else:\n","            denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)\n","\n","        step_size = lr / bias_correction1\n","\n","        param.addcdiv_(exp_avg, denom, value=-step_size)\n","\n","\n","\n","\n","class Adamw(Optimizer):\n","    r\"\"\"Implements AdamW algorithm.\n","\n","    The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_.\n","    The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_.\n","\n","    Args:\n","        params (iterable): iterable of parameters to optimize or dicts defining\n","            parameter groups\n","        lr (float, optional): learning rate (default: 1e-3)\n","        betas (Tuple[float, float], optional): coefficients used for computing\n","            running averages of gradient and its square (default: (0.9, 0.999))\n","        eps (float, optional): term added to the denominator to improve\n","            numerical stability (default: 1e-8)\n","        weight_decay (float, optional): weight decay coefficient (default: 1e-2)\n","        amsgrad (boolean, optional): whether to use the AMSGrad variant of this\n","            algorithm from the paper `On the Convergence of Adam and Beyond`_\n","            (default: False)\n","\n","    .. _Adam\\: A Method for Stochastic Optimization:\n","        https://arxiv.org/abs/1412.6980\n","    .. _Decoupled Weight Decay Regularization:\n","        https://arxiv.org/abs/1711.05101\n","    .. _On the Convergence of Adam and Beyond:\n","        https://openreview.net/forum?id=ryQu7f-RZ\n","    \"\"\"\n","\n","    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,\n","                 weight_decay=1e-2, amsgrad=False):\n","        if not 0.0 <= lr:\n","            raise ValueError(\"Invalid learning rate: {}\".format(lr))\n","        if not 0.0 <= eps:\n","            raise ValueError(\"Invalid epsilon value: {}\".format(eps))\n","        if not 0.0 <= betas[0] < 1.0:\n","            raise ValueError(\"Invalid beta parameter at index 0: {}\".format(betas[0]))\n","        if not 0.0 <= betas[1] < 1.0:\n","            raise ValueError(\"Invalid beta parameter at index 1: {}\".format(betas[1]))\n","        if not 0.0 <= weight_decay:\n","            raise ValueError(\"Invalid weight_decay value: {}\".format(weight_decay))\n","        defaults = dict(lr=lr, betas=betas, eps=eps,\n","                        weight_decay=weight_decay, amsgrad=amsgrad)\n","        super(Adamw, self).__init__(params, defaults)\n","\n","    def __setstate__(self, state):\n","        super(Adamw, self).__setstate__(state)\n","        for group in self.param_groups:\n","            group.setdefault('amsgrad', False)\n","\n","    @torch.no_grad()\n","    def step(self, closure=None):\n","        \"\"\"Performs a single optimization step.\n","\n","        Args:\n","            closure (callable, optional): A closure that reevaluates the model\n","                and returns the loss.\n","        \"\"\"\n","        loss = None\n","        if closure is not None:\n","            with torch.enable_grad():\n","                loss = closure()\n","\n","        for group in self.param_groups:\n","            params_with_grad = []\n","            grads = []\n","            exp_avgs = []\n","            exp_avg_sqs = []\n","            state_sums = []\n","            max_exp_avg_sqs = []\n","            state_steps = []\n","            amsgrad = group['amsgrad']\n","            beta1, beta2 = group['betas']\n","\n","            for p in group['params']:\n","                if p.grad is None:\n","                    continue\n","                params_with_grad.append(p)\n","                if p.grad.is_sparse:\n","                    raise RuntimeError('AdamW does not support sparse gradients')\n","                grads.append(p.grad)\n","\n","                state = self.state[p]\n","\n","                # State initialization\n","                if len(state) == 0:\n","                    state['step'] = 0\n","                    # Exponential moving average of gradient values\n","                    state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n","                    # Exponential moving average of squared gradient values\n","                    state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n","                    if amsgrad:\n","                        # Maintains max of all exp. moving avg. of sq. grad. values\n","                        state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n","\n","                exp_avgs.append(state['exp_avg'])\n","                exp_avg_sqs.append(state['exp_avg_sq'])\n","\n","                if amsgrad:\n","                    max_exp_avg_sqs.append(state['max_exp_avg_sq'])\n","\n","                # update the steps for each param group update\n","                state['step'] += 1\n","                # record the step after step update\n","                state_steps.append(state['step'])\n","\n","            adamw(params_with_grad,\n","                    grads,\n","                    exp_avgs,\n","                    exp_avg_sqs,\n","                    max_exp_avg_sqs,\n","                    state_steps,\n","                    amsgrad=amsgrad,\n","                    beta1=beta1,\n","                    beta2=beta2,\n","                    lr=group['lr'],\n","                    weight_decay=group['weight_decay'],\n","                    eps=group['eps'])\n","\n","        return loss"],"metadata":{"id":"-nHl5bmAk8-L"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# Dataloader\n"],"metadata":{"id":"SjjX6nDgnKwI"}},{"cell_type":"code","source":["# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.\n","\n","# Licensed under the Apache License, Version 2.0 (the \"License\");\n","# you may not use this file except in compliance with the License.\n","# You may obtain a copy of the License at\n","\n","#     http://www.apache.org/licenses/LICENSE-2.0\n","\n","# Unless required by applicable law or agreed to in writing, software\n","# distributed under the License is distributed on an \"AS IS\" BASIS,\n","# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n","# See the License for the specific language governing permissions and\n","# limitations under the License.\n","\n","# Code is heavily based on paper \"Geometry-Informed Neural Operator for Large-Scale 3D PDEs\", we use paddle to reproduce the results of the paper\n","\n","\n","import torch\n","from torch.utils.data import Dataset, DataLoader\n","\n","\n","class BaseDataModule:\n","\n","    @property\n","    def train_dataset(self) -> Dataset:\n","        raise NotImplementedError\n","\n","    @property\n","    def val_dataset(self) -> Dataset:\n","        raise NotImplementedError\n","\n","    @property\n","    def test_dataset(self) -> Dataset:\n","        raise NotImplementedError\n","\n","    def train_dataloader(self, **kwargs) -> DataLoader:\n","        collate_fn = getattr(self, 'collate_fn', None)\n","        return DataLoader(self.train_data, collate_fn=\n","        collate_fn, **kwargs)\n","\n","    def val_dataloader(self, **kwargs) -> DataLoader:\n","        collate_fn = getattr(self, 'collate_fn', None)\n","        return DataLoader(self.val_data, collate_fn=\n","        collate_fn, **kwargs)\n","\n","    def test_dataloader(self, **kwargs) -> DataLoader:\n","        collate_fn = getattr(self, 'collate_fn', None)\n","        return DataLoader(self.test_data, collate_fn=\n","        collate_fn, **kwargs)\n","\n","import torch\n","import open3d as o3d\n","import numpy as np\n","import random\n","import os\n","from pathlib import Path\n","from torch.utils.data import Dataset, DataLoader\n","\n","\n","\n","class AddGaussianNoiseWithProbability(object):\n","    def __init__(self, mean=0.0, std=1.0, probability=0.4):\n","        self.mean = mean\n","        self.std = std\n","        self.probability = probability\n","\n","    def __call__(self, tensor):\n","        if random.random() < self.probability:\n","            noise = torch.randn(tensor.size()) * self.std + self.mean\n","            return tensor + noise\n","        else:\n","            return tensor\n","\n","    def __repr__(self):\n","        return self.__class__.__name__ + '(mean={0}, std={1}, probability={2})'.format(self.mean, self.std,\n","                                                                                       self.probability)\n","\n","class DictDataset(Dataset):\n","    def __init__(self, data_dict: dict):\n","        self.data_dict = data_dict\n","        for k, v in data_dict.items():\n","            assert len(v) == len(\n","                data_dict[list(data_dict.keys())[0]]\n","            ), \"All data must have the same length\"\n","\n","    def __getitem__(self, index):\n","        return {k: v[index] for k, v in self.data_dict.items()}\n","\n","    def __len__(self):\n","        return len(self.data_dict[list(self.data_dict.keys())[0]])\n","\n","\n","class DictDatasetWithConstant(DictDataset):\n","    def __init__(self, data_dict: dict, constant_dict: dict,use_transform=False):\n","        super().__init__(data_dict)\n","        self.constant_dict = constant_dict\n","        self.use_transform=use_transform\n","        self.transform = AddGaussianNoiseWithProbability(0.0, 0.05)\n","    def __getitem__(self, index):\n","        return_dict = {k: v[index] for k, v in self.data_dict.items()}\n","        return_dict.update(self.constant_dict)\n","        if self.use_transform:\n","            return_dict['vertices'] = self.transform(return_dict['vertices'])\n","        return return_dict\n","\n","\n","class CarDataset(BaseDataModule):\n","    def __init__(\n","            self,\n","            train_data_dir,\n","            test_data_dir,\n","            n_train: int = 500,\n","            n_test: int = 5,\n","            val_ratio=0.1,\n","            expand_data=False,\n","            track='track_A',\n","            wind_direction='z',use_transform=False):\n","        BaseDataModule.__init__(self)\n","\n","        valid_train_indices = self.load_valid_mesh_indices(train_data_dir)\n","        valid_test_indices = self.load_valid_mesh_indices(test_data_dir)\n","        random.shuffle(valid_train_indices)\n","        val_num = int(n_train * val_ratio)\n","        valid_val_indices = valid_train_indices[:val_num]\n","        valid_train_indices = valid_train_indices[val_num:]\n","        self.all_train_mesh_path = [self.get_mesh_path_train(train_data_dir, i) for i in valid_train_indices]\n","        self.all_val_mesh_path = [self.get_mesh_path_train(train_data_dir, i) for i in valid_val_indices]\n","        self.all_test_mesh_path = [self.get_mesh_path_test(test_data_dir, i) for i in valid_test_indices]\n","\n","        all_train_vertices = torch.stack(\n","            [\n","                torch.tensor(self.vertices_from_mesh(i)) for i in self.all_train_mesh_path\n","            ]\n","        )\n","        all_val_vertices = torch.stack(\n","            [\n","                torch.tensor(self.vertices_from_mesh(i)) for i in self.all_val_mesh_path\n","            ]\n","        )\n","        all_test_vertices = torch.stack(\n","            [\n","                torch.tensor(self.vertices_from_mesh(i)) for i in self.all_test_mesh_path\n","            ]\n","        )\n","        all_train_press = torch.stack(\n","            [\n","                torch.tensor(\n","                    self.load_pressure(train_data_dir, meshidx), dtype=torch.float32\n","                )\n","                for meshidx in valid_train_indices\n","            ]\n","        )\n","        all_val_press = torch.stack(\n","            [\n","                torch.tensor(\n","                    self.load_pressure(train_data_dir, meshidx), dtype=torch.float32\n","                )\n","                for meshidx in valid_val_indices\n","            ]\n","        )\n","        all_test_press = torch.stack(\n","            [\n","                torch.tensor(\n","                    self.load_pressure(test_data_dir, meshidx), dtype=torch.float32\n","                )\n","                for meshidx in valid_test_indices\n","            ]\n","        )\n","        if track == 'track_A':\n","            track_token = torch.tensor([1, 0, 0], dtype=torch.float32).unsqueeze(0)\n","        else:\n","            track_token = torch.tensor([0, 1, 0], dtype=torch.float32).unsqueeze(0)\n","        if wind_direction == 'x':\n","            wind_token = torch.tensor([1, 0, 0], dtype=torch.float32).unsqueeze(0)\n","        elif wind_direction == 'z':\n","            wind_token = torch.tensor([0, 1, 0], dtype=torch.float32).unsqueeze(0)\n","        if expand_data:\n","            all_train_vertices = torch.cat([all_train_vertices, all_train_vertices], dim=0)\n","            all_train_press = torch.cat([all_train_press, all_train_press], dim=0)\n","        print(f'use {len(all_train_press)} training data!')\n","        self._train_data = DictDatasetWithConstant(\n","            {\"vertices\": all_train_vertices,\n","             \"pressure\": all_train_press},\n","            {\"track\": track_token,\n","             'wind_direction': wind_token},\n","            use_transform\n","        )\n","        self._val_data = DictDatasetWithConstant(\n","            {\"vertices\": all_val_vertices,\n","             \"pressure\": all_val_press},\n","            {\"track\": track_token,\n","             'wind_direction': wind_token},\n","        )\n","        self._test_data = DictDatasetWithConstant(\n","            {\"vertices\": all_test_vertices,\n","             \"pressure\": all_test_press},\n","            {\"track\": track_token,\n","             'wind_direction': wind_token},\n","        )\n","\n","    @property\n","    def train_data(self):\n","        return self._train_data\n","\n","    @property\n","    def test_data(self):\n","        return self._test_data\n","\n","    @property\n","    def val_data(self):\n","        return self._val_data\n","\n","    def load_pressure(self, data_dir: Path, mesh_index: int) -> np.ndarray:\n","        press_path = self.get_pressure_data_path(data_dir, mesh_index)\n","        assert press_path.exists(), \"Pressure data does not exist\"\n","        press = np.load(press_path).reshape((-1,)).astype(np.float32)\n","        press = np.concatenate((press[0:16], press[112:]), axis=0)\n","        return press\n","\n","    def get_mesh_path(self, data_dir: Path, mesh_ind: int) -> Path:\n","        return data_dir / (\"mesh_\" + str(mesh_ind).zfill(3) + \".ply\")\n","\n","    def get_mesh_path_train(self, data_dir: Path, mesh_ind: int) -> Path:\n","        return data_dir / 'Feature' / (\"mesh_\" + str(mesh_ind).zfill(3) + \".ply\")\n","\n","    def get_mesh_path_test(self, data_dir: Path, mesh_ind: int) -> Path:\n","        return data_dir / 'Inference' / (\"mesh_\" + str(mesh_ind).zfill(3) + \".ply\")\n","\n","    def get_pressure_data_path(self, data_dir: Path, mesh_ind: int) -> Path:\n","        return data_dir / 'Label' / (\"press_\" + str(mesh_ind).zfill(3) + \".npy\")\n","\n","    def load_mesh(self, mesh_path: Path) -> o3d.t.geometry.TriangleMesh:\n","        assert mesh_path.exists(), \"Mesh path does not exist\"\n","        mesh = o3d.io.read_triangle_mesh(str(mesh_path))\n","        mesh = o3d.t.geometry.TriangleMesh.from_legacy(mesh)\n","        return mesh\n","\n","    def vertices_from_mesh(self, mesh_path: Path):\n","        mesh = self.load_mesh(mesh_path)\n","        vertices = mesh.vertex.positions.numpy()\n","        return vertices\n","\n","    def load_valid_mesh_indices(\n","            self, data_dir, filename=\"watertight_meshes.txt\"\n","    ):\n","        with open(data_dir / filename, \"r\") as fp:\n","            mesh_ind = fp.read().split(\"\\n\")\n","            mesh_ind = [int(a) for a in mesh_ind if a.strip()]\n","        return mesh_ind\n","\n","\n","class CarDataset_trackB(Dataset):\n","    def __init__(\n","            self,\n","            train_mesh_dir,\n","            train_centroid_dir,\n","            test_mesh_dir,\n","            n_train: int = 500,\n","            n_test: int = 5,\n","            val_ratio=0.1,\n","            mode='train',\n","            track='track_B',\n","            wind_direction='x',\n","            valid_val_indices=None):\n","        Dataset.__init__(self)\n","\n","        valid_train_indices = self.get_all_mesh_indices(train_mesh_dir)\n","        # valid_test_indices = self.get_all_mesh_indices(test_mesh_dir)\n","        random.shuffle(valid_train_indices)\n","        val_num = int(n_train * val_ratio)\n","        if valid_val_indices is not None:\n","            valid_val_indices = valid_val_indices\n","        else:\n","            valid_val_indices = valid_train_indices[:val_num]\n","        valid_train_indices = valid_train_indices[val_num:]\n","        self.mode = mode\n","        self.valid_val_indices = valid_val_indices\n","        self.all_train_mesh_path = [self.get_mesh_path_train(train_mesh_dir, i) for i in valid_train_indices]\n","        self.all_train_press_path = [self.get_pressure_data_path(train_centroid_dir, i) for i in valid_train_indices]\n","        self.all_val_mesh_path = [self.get_mesh_path_train(train_mesh_dir, i) for i in valid_val_indices]\n","        self.all_val_press_path = [self.get_pressure_data_path(train_centroid_dir, i) for i in valid_val_indices]\n","        self.all_test_mesh_path = [self.get_mesh_path_test(test_mesh_dir, i) for i in range(1, 51)]\n","        self.all_train_centroid_path = [self.get_centroid_data_path(train_centroid_dir, i) for i in valid_train_indices]\n","        self.all_val_centroid_path = [self.get_centroid_data_path(train_centroid_dir, i) for i in valid_val_indices]\n","        self.all_test_centroid_path = [self.get_centroid_data_path_test(test_mesh_dir, i) for i in range(1, 51)]\n","        self.all_train_area_path = [self.get_area_data_path(train_centroid_dir, i) for i in valid_train_indices]\n","        self.all_val_area_path = [self.get_area_data_path(train_centroid_dir, i) for i in valid_val_indices]\n","        self.all_test_area_path = [self.get_area_data_path_test(test_mesh_dir, i) for i in range(1, 51)]\n","\n","        if track == 'track_A':\n","            self.track_token = torch.tensor([1, 0, 0], dtype=torch.float32).unsqueeze(0)\n","        else:\n","            self.track_token = torch.tensor([0, 1, 0], dtype=torch.float32).unsqueeze(0)\n","        if wind_direction == 'x':\n","            self.wind_token = torch.tensor([1, 0, 0], dtype=torch.float32).unsqueeze(0)\n","        elif wind_direction == 'z':\n","            self.wind_token = torch.tensor([0, 1, 0], dtype=torch.float32).unsqueeze(0)\n","\n","        print(f'use {len(self.all_train_mesh_path)} training data!')\n","\n","    def __getitem__(self, idx):\n","        if self.mode == 'train':\n","            data_dict = {\n","                'vertices': self.load_centroid(self.all_train_centroid_path[idx]),\n","                'pressure': self.load_pressure(self.all_train_press_path[idx]),\n","                'area': self.load_area(self.all_train_area_path[idx]),\n","                'track': self.track_token,\n","                'wind_direction': self.wind_token,\n","\n","            }\n","        elif self.mode == 'val':\n","            data_dict = {\n","                'vertices': self.load_centroid(self.all_val_centroid_path[idx]),\n","                'pressure': self.load_pressure(self.all_val_press_path[idx]),\n","                'area': self.load_area(self.all_val_area_path[idx]),\n","                'track': self.track_token,\n","                'wind_direction': self.wind_token,\n","\n","            }\n","        else:\n","            data_dict = {\n","                'vertices': self.load_centroid(self.all_test_centroid_path[idx]),\n","                # 'pressure':self.load_pressure(self.all_te_press_path[idx]),\n","                'area': self.load_area(self.all_test_area_path[idx]),\n","                'track': self.track_token,\n","                'wind_direction': self.wind_token,\n","\n","            }\n","        return data_dict\n","\n","    def __len__(self):\n","        if self.mode == 'train':\n","            return len(self.all_train_mesh_path)\n","        elif self.mode == 'val':\n","            return len(self.all_val_mesh_path)\n","        else:\n","            return len(self.all_test_mesh_path)\n","\n","    def load_pressure(self, data_path: Path) -> np.ndarray:\n","        # press_path = self.get_pressure_data_path(data_dir, mesh_index)\n","        assert data_path.exists(), \"Pressure data does not exist\"\n","        press = np.load(data_path).reshape((-1,)).astype(np.float32)\n","        # press = np.concatenate((press[0:16], press[112:]), axis=0)\n","        return press\n","\n","    def load_centroid(self, data_path: Path) -> np.ndarray:\n","        assert data_path.exists(), \"Pressure data does not exist\"\n","        centroid = np.load(data_path).astype(np.float32)\n","        return centroid\n","\n","    def load_area(self, data_path):\n","        assert data_path.exists(), \"area data does not exist\"\n","        area = np.load(data_path).reshape((-1,)).astype(np.float32)\n","\n","        return area\n","\n","    def get_mesh_path(self, data_dir: Path, mesh_ind: int) -> Path:\n","        return data_dir / (\"mesh_\" + str(mesh_ind).zfill(4) + \".ply\")\n","\n","    def get_mesh_path_train(self, data_dir: Path, mesh_ind: int) -> Path:\n","        return data_dir / 'train_mesh_0603' / (\"mesh_\" + str(mesh_ind).zfill(4) + \".ply\")\n","\n","    def get_mesh_path_test(self, data_dir: Path, mesh_ind: int) -> Path:\n","        return data_dir / 'Inference' / (\"mesh_\" + str(mesh_ind) + \".ply\")\n","\n","    def get_pressure_data_path(self, data_dir: Path, mesh_ind: int) -> Path:\n","        return data_dir / (\"press_\" + str(mesh_ind).zfill(4) + \".npy\")\n","\n","    def get_area_data_path(self, data_dir: Path, mesh_ind: int) -> Path:\n","        return data_dir / (\"area_\" + str(mesh_ind).zfill(4) + \".npy\")\n","\n","    def get_area_data_path_test(self, data_dir: Path, mesh_ind: int) -> Path:\n","        return data_dir / 'Auxiliary' / (\"area_\" + str(mesh_ind) + \".npy\")\n","\n","    def get_centroid_data_path(self, data_dir: Path, mesh_ind: int) -> Path:\n","        return data_dir / (\"centroid_\" + str(mesh_ind).zfill(4) + \".npy\")\n","\n","    def get_centroid_data_path_test(self, data_dir: Path, mesh_ind: int) -> Path:\n","        return data_dir / 'Inference' / (\"centroid_\" + str(mesh_ind) + \".npy\")\n","\n","    def load_mesh(self, mesh_path: Path) -> o3d.t.geometry.TriangleMesh:\n","        assert mesh_path.exists(), \"Mesh path does not exist\"\n","        mesh = o3d.io.read_triangle_mesh(str(mesh_path))\n","        mesh = o3d.t.geometry.TriangleMesh.from_legacy(mesh)\n","        return mesh\n","\n","    def vertices_from_mesh(self, mesh_path: Path):\n","        mesh = self.load_mesh(mesh_path)\n","        vertices = mesh.vertex.positions.numpy()\n","        return vertices\n","\n","    def get_all_mesh_indices(\n","            self, mesh_data_dir\n","    ):\n","        mesh_ind = []\n","        all_mesh_files = os.listdir(str(mesh_data_dir))\n","        all_mesh_files.sort()\n","        for mesh_file_name in all_mesh_files:\n","            mesh_idx = mesh_file_name.split('_')[1].split('.')[0]\n","            mesh_ind.append(mesh_idx)\n","        return mesh_ind\n","\n","    def get_valid_val_indices(self):\n","        return self.valid_val_indices\n","\n","\n","\n"],"metadata":{"id":"xVtHRKMfkbhK"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["#Test"],"metadata":{"id":"Ktwbrdh7m_gv"}},{"cell_type":"markdown","source":["wget导入权重"],"metadata":{"id":"9mssn11Uyhku"}},{"cell_type":"code","source":["!wget  --no-check-certificate 'https://docs.google.com/uc?export=download&id=1lscczzYFotpzOPPoMlZq4XP_Dv2mHI-1' -O best_trackb.pth"],"metadata":{"id":"XdIpBndZrRLG"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["#!/usr/bin/env python\n","# -*- coding:utf-8 _*-\n","import sys\n","import os\n","\n","sys.path.append('../..')\n","sys.path.append('..')\n","\n","import re\n","import time\n","import pickle\n","import numpy as np\n","import torch\n","import torch.nn as nn\n","\n","\n","from torch.utils.data import DataLoader\n","from pathlib import Path\n","\n","from tqdm import tqdm\n","\n","\n","'''\n","    A general code framework for training neural operator on irregular domains\n","'''\n","\n","EPOCH_SCHEDULERS = ['ReduceLROnPlateau', 'StepLR', 'MultiplicativeLR',\n","                    'MultiStepLR', 'ExponentialLR', 'LambdaLR']\n","\n","\n","def extract_numbers(s):\n","    return [int(digit) for digit in re.findall(r'\\d+', s)]\n","\n","\n","def write_to_vtk(p, point_data_pos=\"press on mesh cells\", mesh_path=None):\n","    import meshio\n","    # p = out_dict[\"pressure\"]\n","    index = extract_numbers(mesh_path.name)[0]\n","    index = str(index)\n","\n","    if point_data_pos == \"press on mesh points\":\n","        mesh = meshio.read(mesh_path)\n","        mesh.point_data[\"p\"] = p.cpu().numpy()\n","        # if \"pred wss_x\" in out_dict:\n","        #     wss_x = out_dict[\"pred wss_x\"]\n","        #     mesh.point_data[\"wss_x\"] = wss_x.numpy()\n","    elif point_data_pos == \"press on mesh cells\":\n","        points = np.load(mesh_path.parent / f\"centroid_{index}.npy\")\n","        npoint = points.shape[0]\n","        mesh = meshio.Mesh(\n","            points=points, cells=[(\"vertex\", np.arange(npoint).reshape(npoint, 1))]\n","        )\n","        mesh.point_data = {\"p\": p.cpu().numpy()}\n","    pressure = p.cpu().numpy().flatten()\n","    # print(f\"write : {config.run_name}/output/{mesh_path.parent.name}_{index}.vtk\")\n","    # os.makedirs(f\"./result/{args.track}/output\", exist_ok=True)\n","    os.makedirs(f\"./result\", exist_ok=True)\n","    np.save(f\"./result/press_{index}.npy\", pressure)\n","    # mesh.write(f\"./result/{args.track}/output/{mesh_path.parent.name}_{index}.vtk\")\n","\n","\n","\n","\n","\n","def test_epoch(model, metric_func, test_loader, device, model_weight):\n","    model.load_state_dict(model_weight)\n","    model.eval()\n","    metric_val = []\n","    for i, data in enumerate(test_loader):\n","        with torch.no_grad():\n","            query_points,  area = data['vertices'], data['area'].unsqueeze(-1)\n","            query_points = query_points.to(device)\n","            conditions = extract_condition(data)\n","            if args.cat_area:\n","                area = area.to(device)\n","                query_points = torch.cat([query_points, area], dim=-1)\n","            else:\n","                conditions.append(area)\n","            conditions = MultipleTensors(conditions)\n","            conditions = conditions.to(device)\n","\n","            out = model(query_points, conditions)\n","\n","            y_pred = out.squeeze(-1)\n","            write_to_vtk(y_pred.unsqueeze(-1).squeeze(0), mesh_path=CARdataset_test.all_test_mesh_path[i])\n","            # metric = metric_func(y_pred, y)\n","\n","            # metric_val.append(metric.cpu().numpy())\n","\n","    # return dict(metric=np.mean(metric_val, axis=0))\n","\n","\n","if __name__ == \"__main__\":\n","    args = get_args()\n","\n","    if not args.no_cuda and torch.cuda.is_available():\n","        device = torch.device('cuda:{}'.format(str(args.gpu)))\n","    else:\n","        device = torch.device(\"cpu\")\n","\n","\n","    kwargs = {'pin_memory': False} if args.gpu else {}\n","    get_seed(args.seed, printout=False)\n","\n","    train_dir = Path('./Dataset/train_mesh_0603')\n","    test_dir = Path('./Dataset/track_B')\n","    train_cent_dir = Path('./Dataset/train_track_B')\n","    test_aux_dir = Path('./Dataset/track_B/Auxiliary')\n","\n","    CARdataset_test = CarDataset_trackB(train_dir, train_cent_dir, test_dir, mode='test')\n","\n","\n","    test_loader = DataLoader(CARdataset_test,\n","                             batch_size=1, shuffle=False\n","                             )\n","\n","    # args.space_dim = int(re.search(r'\\d', args.dataset).group())\n","    # args.normalizer =  train_dataset.y_normalizer.to(device) if train_dataset.y_normalizer is not None else None\n","\n","    #### set random seeds\n","    get_seed(args.seed)\n","    torch.cuda.empty_cache()\n","\n","    loss_func = LpLoss()\n","    metric_func = LpLoss()\n","\n","    model = get_model(args)\n","    model = model.to(device)\n","    print(f\"\\nModel: {model.__name__}\\t Number of params: {get_num_params(model)}\")\n","\n","\n","\n","    time_start = time.time()\n","\n","    print('Training takes {} seconds.'.format(time.time() - time_start))\n","\n","    best_state_dict = torch.load('./best_trackb.pth',map_location='cuda:0')\n","    print(f\"\\nBest model's validation metric in this run: \")\n","\n","    test_epoch(model, metric_func, test_loader, device, best_state_dict)\n","\n","\n"],"metadata":{"id":"oMfabw83kmqT"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# !sh test_tracka.sh\n","!zip -r ./B_result.zip ./result/*"],"metadata":{"collapsed":true,"id":"5hq2LJUH69Ra"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# Train!\n"],"metadata":{"id":"Ce0Ef-gOobJn"}},{"cell_type":"markdown","source":["训练结束后可以在 runs/train/output_pressure下得到输出结果"],"metadata":{"id":"P51Ja8cw0pgk"}},{"cell_type":"code","source":["#!/usr/bin/env python\n","# -*- coding:utf-8 _*-\n","import sys\n","import os\n","\n","sys.path.append('../..')\n","sys.path.append('..')\n","\n","import re\n","import time\n","import pickle\n","import numpy as np\n","import torch\n","import torch.nn as nn\n","import warnings\n","from torch.optim.lr_scheduler import OneCycleLR, StepLR, LambdaLR\n","from torch.utils.tensorboard import SummaryWriter\n","from torch.utils.data import DataLoader\n","from pathlib import Path\n","from tqdm import tqdm\n","from torch.cuda.amp import GradScaler, autocast\n","\n","'''\n","    A general code framework for training neural operator on irregular domains\n","'''\n","\n","EPOCH_SCHEDULERS = ['ReduceLROnPlateau', 'StepLR', 'MultiplicativeLR',\n","                    'MultiStepLR', 'ExponentialLR', 'LambdaLR']\n","\n","\n","def extract_numbers(s):\n","    return [int(digit) for digit in re.findall(r'\\d+', s)]\n","\n","\n","def write_to_vtk(p, point_data_pos=\"press on mesh cells\", mesh_path=None):\n","    import meshio\n","    # p = out_dict[\"pressure\"]\n","    index = extract_numbers(mesh_path.name)[0]\n","    index = str(index)\n","\n","    if point_data_pos == \"press on mesh points\":\n","        mesh = meshio.read(mesh_path)\n","        mesh.point_data[\"p\"] = p.cpu().numpy()\n","        # if \"pred wss_x\" in out_dict:\n","        #     wss_x = out_dict[\"pred wss_x\"]\n","        #     mesh.point_data[\"wss_x\"] = wss_x.numpy()\n","    elif point_data_pos == \"press on mesh cells\":\n","        points = np.load(mesh_path.parent / f\"centroid_{index}.npy\")\n","        npoint = points.shape[0]\n","        mesh = meshio.Mesh(\n","            points=points, cells=[(\"vertex\", np.arange(npoint).reshape(npoint, 1))]\n","        )\n","        mesh.point_data = {\"p\": p.cpu().numpy()}\n","    pressure = p.cpu().numpy().flatten()\n","    # print(f\"write : {config.run_name}/output/{mesh_path.parent.name}_{index}.vtk\")\n","    os.makedirs(f\"./runs/{args.train_log}/output\", exist_ok=True)\n","    os.makedirs(f\"./runs/{args.train_log}/output_pressure\", exist_ok=True)\n","    np.save(f\"./runs/{args.train_log}/output_pressure/press_{index}.npy\", pressure)\n","    mesh.write(f\"./runs/{args.train_log}/output/{mesh_path.parent.name}_{index}.vtk\")\n","\n","\n","def train(model, loss_func, metric_func,\n","          train_loader, valid_loader,\n","          optimizer, lr_scheduler,\n","          epochs=10,\n","          writer=None,\n","          device=\"cuda\",\n","          patience=10,\n","          grad_clip=0.999,\n","          start_epoch: int = 0,\n","          model_save_path='./data/checkpoints/',\n","          save_mode='state_dict',  # 'state_dict' or 'entire'\n","          ):\n","    loss_train = []\n","    loss_val = []\n","    loss_epoch = []\n","    lr_history = []\n","    it = 0\n","\n","    if patience is None or patience == 0:\n","        patience = epochs\n","    result = None\n","    start_epoch = start_epoch\n","    end_epoch = start_epoch + epochs\n","    best_val_metric = np.inf\n","    best_val_epoch = None\n","    save_mode = 'state_dict' if save_mode is None else save_mode\n","    stop_counter = 0\n","    is_epoch_scheduler = any(s in str(lr_scheduler.__class__) for s in EPOCH_SCHEDULERS)\n","\n","    for epoch in range(start_epoch, end_epoch):\n","        model.train()\n","        torch.cuda.empty_cache()\n","        for batch in tqdm(train_loader):\n","            loss_epoch = []\n","            loss = train_batch(model, loss_func, batch, optimizer, lr_scheduler, device, grad_clip=grad_clip)\n","            loss = np.array(loss)\n","            loss_epoch.append(loss)\n","            it += 1\n","            lr = optimizer.param_groups[0]['lr']\n","            writer.add_scalar('LR', lr, it)\n","            lr_history.append(lr)\n","            # log = f\"epoch: [{epoch + 1}/{end_epoch}]\"\n","            # if loss.ndim == 0:  # 1 target loss\n","            #     _loss_mean = np.mean(loss_epoch)\n","            #     log += \" loss: {:.6f}\".format(_loss_mean)\n","            # else:\n","        _loss_mean = np.mean(loss_epoch)\n","        #     for j in range(len(_loss_mean)):\n","        #         log += \" | loss {}: {:.6f}\".format(j, _loss_mean[j])\n","        # log += \" | current lr: {:.3e}\".format(lr)\n","        #\n","        # if it % print_freq == 0:\n","        #     print(log)\n","\n","        if writer is not None:\n","            writer.add_scalar(\"train_loss\", _loss_mean, epoch)  #### loss 0 seems to be the sum of all loss\n","\n","        loss_train.append(_loss_mean)\n","        if epoch % args.val_epochs==0:\n","            val_result = validate_epoch(model, metric_func, valid_loader, device)\n","\n","            loss_val.append(val_result[\"metric\"])\n","            val_metric = val_result[\"metric\"]\n","            print(\n","                f'epoch:[{epoch + 1}/{end_epoch}] train_loss:{_loss_mean:.4f} val_loss:{val_metric:.4f} lr:{lr_history[-1]:.6f}')\n","        else:\n","            print(\n","                f'epoch:[{epoch + 1}/{end_epoch}] train_loss:{_loss_mean:.4f} lr:{lr_history[-1]:.6f}')\n","\n","        if val_metric < best_val_metric:\n","            best_val_epoch = epoch\n","            best_val_metric = val_metric\n","            torch.save(model.state_dict(), os.path.join(model_save_path, 'best_val.pth'))\n","\n","        if lr_scheduler and is_epoch_scheduler:\n","            if 'ReduceLROnPlateau' in str(lr_scheduler.__class__):\n","                lr_scheduler.step(val_metric)\n","            else:\n","                lr_scheduler.step()\n","\n","        # if val_result[\"metric\"].size == 1:\n","        #     log = \"| val metric 0: {:.6f} \".format(val_metric)\n","        #\n","        # else:\n","        #     log = ''\n","        #     for i, metric_i in enumerate(val_result['metric']):\n","        #         log += '| val metric {} : {:.6f} '.format(i, metric_i)\n","\n","        if writer is not None:\n","            writer.add_scalar('val loss', val_metric, epoch)\n","\n","        # log += \"| best val: {:.6f} at epoch {} | current lr: {:.3e}\".format(best_val_metric, best_val_epoch + 1, lr)\n","\n","        # desc_ep = \"\"\n","        # if _loss_mean.ndim == 0:  # 1 target loss\n","        #     desc_ep += \"| loss: {:.6f}\".format(_loss_mean)\n","        # else:\n","        #     for j in range(len(_loss_mean)):\n","        #         if _loss_mean[j] > 0:\n","        #             desc_ep += \"| loss {}: {:.3e}\".format(j, _loss_mean[j])\n","        #\n","        # desc_ep += log\n","        # print(desc_ep)\n","        torch.save(model.state_dict(), os.path.join(f'runs/{args.train_log}/last.pth'))\n","        result = dict(\n","            # best_val_epoch=best_val_epoch,\n","            # best_val_metric=best_val_metric,\n","            # loss_train=np.asarray(loss_train),\n","            # loss_val=np.asarray(loss_val),\n","            # lr_history=np.asarray(lr_history),\n","            # best_model=best_model_state_dict,\n","            # optimizer_state=optimizer.state_dict()\n","        )\n","        # pickle.dump(result, open(os.path.join(model_save_path, result_name), 'wb'))\n","    return result\n","\n","\n","def train_batch(model, loss_func, data, optimizer, lr_scheduler, device, grad_clip=0.999):\n","    # scaler = GradScaler()\n","    optimizer.zero_grad()\n","\n","    query_points, y, area = data['vertices'], data['pressure'], data['area'].unsqueeze(-1)\n","    query_points, y = query_points.to(device), y.to(device),\n","    conditions = extract_condition(data)\n","    if args.cat_area:\n","        area = area.to(device)\n","        query_points = torch.cat([query_points, area], dim=-1)\n","    else:\n","        conditions.append(area)\n","    conditions = MultipleTensors(conditions)\n","    conditions = conditions.to(device)\n","\n","    out = model(query_points, conditions)\n","\n","    y_pred = out.squeeze(-1)\n","    loss = loss_func(y_pred, y)\n","\n","    # scaler.scale(loss).backward()\n","    # scaler.step(optimizer)\n","    # scaler.update()\n","\n","    loss.backward()\n","    nn.utils.clip_grad_norm_(model.parameters(), grad_clip)\n","    optimizer.step()\n","\n","    if lr_scheduler:\n","        lr_scheduler.step()\n","\n","    return loss.item()\n","\n","\n","def validate_epoch(model, metric_func, valid_loader, device):\n","    model.eval()\n","    metric_val = []\n","    for i, data in enumerate(valid_loader):\n","        with torch.no_grad():\n","            query_points, y, area = data['vertices'], data['pressure'], data['area'].unsqueeze(-1)\n","            query_points, y = query_points.to(device), y.to(device),\n","            conditions = extract_condition(data)\n","            if args.cat_area:\n","                area = area.to(device)\n","                query_points = torch.cat([query_points, area], dim=-1)\n","            else:\n","                conditions.append(area)\n","            conditions = MultipleTensors(conditions)\n","            conditions = conditions.to(device)\n","\n","            out = model(query_points, conditions)\n","\n","            y_pred = out.squeeze(-1)\n","            metric = metric_func(y_pred, y)\n","\n","            metric_val.append(metric.cpu().numpy())\n","\n","    return dict(metric=np.mean(metric_val, axis=0))\n","\n","\n","def test_epoch(model, metric_func, test_loader, device, model_weight):\n","    model.load_state_dict(model_weight)\n","    model.eval()\n","    metric_val = []\n","    for i, data in enumerate(test_loader):\n","        with torch.no_grad():\n","            query_points, area = data['vertices'], data['area'].unsqueeze(-1)\n","            query_points = query_points.to(device)\n","            conditions = extract_condition(data)\n","            if args.cat_area:\n","                area = area.to(device)\n","                query_points = torch.cat([query_points, area], dim=-1)\n","            else:\n","                conditions.append(area)\n","            conditions = MultipleTensors(conditions)\n","            conditions = conditions.to(device)\n","\n","            out = model(query_points, conditions)\n","\n","            y_pred = out.squeeze(-1)\n","            write_to_vtk(y_pred.unsqueeze(-1).squeeze(0), mesh_path=CARdataset_test.all_test_mesh_path[i])\n","            # metric = metric_func(y_pred, y)\n","\n","            # metric_val.append(metric.cpu().numpy())\n","\n","    # return dict(metric=np.mean(metric_val, axis=0))\n","\n","\n","if __name__ == \"__main__\":\n","    warnings.filterwarnings('ignore')\n","    args = get_args()\n","    print(args)\n","    if args.use_tb:\n","        writer = SummaryWriter(f'runs/{args.train_log}')\n","\n","    else:\n","        writer = None\n","        log_path = None\n","    if not args.no_cuda and torch.cuda.is_available():\n","        device = torch.device('cuda:{}'.format(str(args.gpu)))\n","    else:\n","        device = torch.device(\"cpu\")\n","    # 获取命令行参数\n","    arg = sys.argv[1:]\n","\n","    # 将参数写入文件\n","    with open(f'runs/{args.train_log}/args.txt', 'w') as f:\n","        for a in arg:\n","            f.write(a + '\\n')\n","\n","    kwargs = {'pin_memory': False} if args.gpu else {}\n","    get_seed(args.seed, printout=False)\n","\n","    train_dir = Path('./Dataset/train_mesh_0603')\n","    test_dir = Path('./Dataset/track_B')\n","    train_cent_dir = Path('./Dataset/train_track_B')\n","    test_aux_dir = Path('./Dataset/track_B/Auxiliary')\n","\n","    CARdataset_train = CarDataset_trackB(train_dir, train_cent_dir, test_dir, val_ratio=args.val_ratio)\n","    valid_val_indices = CARdataset_train.get_valid_val_indices()\n","    CARdataset_val = CarDataset_trackB(train_dir, train_cent_dir, test_dir, mode='val', val_ratio=args.val_ratio,\n","                                       valid_val_indices=valid_val_indices)\n","    CARdataset_test = CarDataset_trackB(train_dir, train_cent_dir, test_dir, mode='test')\n","    # test_dataset = get_dataset(args)\n","\n","    train_loader = DataLoader(CARdataset_train,\n","                              batch_size=args.batch_size, shuffle=True\n","                              )\n","    val_loader = DataLoader(CARdataset_val,\n","                            batch_size=1, shuffle=False\n","                            )\n","    test_loader = DataLoader(CARdataset_test,\n","                             batch_size=1, shuffle=False\n","                             )\n","\n","    # args.space_dim = int(re.search(r'\\d', args.dataset).group())\n","    # args.normalizer =  train_dataset.y_normalizer.to(device) if train_dataset.y_normalizer is not None else None\n","\n","    #### set random seeds\n","    get_seed(args.seed)\n","    torch.cuda.empty_cache()\n","\n","    loss_func = LpLoss()\n","    metric_func = LpLoss()\n","\n","    model = get_model(args)\n","    model = model.to(device)\n","    print(f\"\\nModel: {model.__name__}\\t Number of params: {get_num_params(model)}\")\n","\n","    # path_prefix = args.dataset + '_{}_'.format(args.component) + model.__name__ + args.comment + time.strftime(\n","    #     '_%m%d_%H_%M_%S')\n","    # model_path, result_path = path_prefix + '.pt', path_prefix + '.pkl'\n","\n","    # print(f\"Saving model and result in ./../models/checkpoints/{model_path}\\n\")\n","\n","    # print(model)\n","    # print(config)\n","\n","    epochs = args.epochs\n","    lr = args.lr\n","\n","    if args.optimizer == 'Adam':\n","        optimizer = Adam(model.parameters(), lr=lr, weight_decay=args.weight_decay, betas=(0.9, 0.999))\n","    elif args.optimizer == \"AdamW\":\n","        optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=args.weight_decay, betas=(0.9, 0.999))\n","    else:\n","        raise NotImplementedError\n","\n","    if args.lr_method == 'cycle':\n","        print('Using cycle learning rate schedule')\n","        scheduler = OneCycleLR(optimizer, max_lr=lr, div_factor=1e4, pct_start=0.2, final_div_factor=1e4,\n","                               steps_per_epoch=len(train_loader), epochs=epochs)\n","    elif args.lr_method == 'step':\n","        print('Using step learning rate schedule')\n","        scheduler = StepLR(optimizer, step_size=args.lr_step_size * len(train_loader), gamma=0.5)\n","    elif args.lr_method == 'warmup':\n","        print('Using warmup learning rate schedule')\n","        scheduler = LambdaLR(optimizer, lambda steps: min((steps + 1) / (args.warmup_epochs * len(train_loader)),\n","                                                          np.power(\n","                                                              args.warmup_epochs * len(train_loader) / float(steps + 1),\n","                                                              0.5)))\n","\n","    time_start = time.time()\n","\n","    result = train(model, loss_func, metric_func,\n","                   train_loader, val_loader,\n","                   optimizer, scheduler,\n","                   epochs=epochs,\n","                   grad_clip=args.grad_clip,\n","                   patience=10,\n","                   model_save_path=f'runs/{args.train_log}',\n","                   writer=writer,\n","                   device=device)\n","\n","    print('Training takes {} seconds.'.format(time.time() - time_start))\n","\n","    # result['args'], result['config'] = args, config\n","    # checkpoint = {'args': args, 'model': model.state_dict(), 'optimizer': optimizer.state_dict()}\n","    torch.save(model.state_dict(), os.path.join(f'runs/{args.train_log}/last.pth'))\n","    # model.eval()\n","    best_state_dict = torch.load(os.path.join(f'runs/{args.train_log}', 'best_val.pth'))\n","    last_state_dict = torch.load(os.path.join(f'runs/{args.train_log}', 'last.pth'))\n","\n","    test_epoch(model, metric_func, test_loader, device, best_state_dict)\n"],"metadata":{"id":"85tvzBxnocZH"},"execution_count":null,"outputs":[]}]}